diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index 3a54a3230..e0c2968f7 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -18,7 +18,10 @@ for transform in (:permute, :transpose) # and store that instead of repeating the permutation in the pullback each time. # effectively, we replace `add_permute` by `add ∘ permute`. Ap = if !isa(α, Const) - Ap = $transform(A.val, p.val) + # this may just alias A if the permutation is an identity or the space + # is Trivial and has_shared_permute is true. + Ap_needs_copy = (p.val == (TensorKit.codomainind(A.val), TensorKit.domainind(A.val)) || TensorKit.has_shared_permute(A.val, p.val)) + Ap = $transform(A.val, p.val; copy = Ap_needs_copy) add!(C.val, Ap, α.val, β.val) Ap else diff --git a/test/enzyme/inv.jl b/test/enzyme/inv.jl index 35bd074fb..20475e4e9 100644 --- a/test/enzyme/inv.jl +++ b/test/enzyme/inv.jl @@ -21,12 +21,9 @@ TDs = is_ci ? (Duplicated,) : (Const, Duplicated) D2 = randn(T, V[1] ⊗ V[2] ← V[1] ⊗ V[2]) EnzymeTestUtils.test_reverse(inv, TD, (D2, TD); atol, rtol) EnzymeTestUtils.test_forward(inv, TD, (D2, TD); atol, rtol) - # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 - @static if VERSION ≥ v"1.11.0-rc" - D3 = randn(T, V[1] ⊗ V[2] ⊗ V[3] ← V[1] ⊗ V[2] ⊗ V[3]) - EnzymeTestUtils.test_reverse(inv, TD, (D3, TD); atol, rtol) - EnzymeTestUtils.test_forward(inv, TD, (D3, TD); atol, rtol) - end + D3 = randn(T, V[1] ⊗ V[2] ⊗ V[3] ← V[1] ⊗ V[2] ⊗ V[3]) + EnzymeTestUtils.test_reverse(inv, TD, (D3, TD); atol, rtol) + EnzymeTestUtils.test_forward(inv, TD, (D3, TD); atol, rtol) end end end diff --git a/test/enzyme/norm.jl b/test/enzyme/norm.jl index 504e51884..d12dccc1a 100644 --- a/test/enzyme/norm.jl +++ b/test/enzyme/norm.jl @@ -14,12 +14,7 @@ fRTs = is_ci ? (Duplicated,) : (Const, Duplicated) @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T), TC $TC" for V in spacelist, T in eltypes, TC in (Const, Duplicated) atol = default_tol(T) rtol = default_tol(T) - # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 - @static if VERSION < v"1.11.0-rc" - C = randn(T, V[1] ⊗ V[2] ← (V[4] ⊗ V[5])') - else - C = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') - end + C = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') for RT in rRTs EnzymeTestUtils.test_reverse(norm, RT, (C, TC), (2, Const); atol, rtol) EnzymeTestUtils.test_reverse(norm, RT, (C', TC), (2, Const); atol, rtol) diff --git a/test/enzyme/permute.jl b/test/enzyme/permute.jl index 4e898f3a3..ffe60f89f 100644 --- a/test/enzyme/permute.jl +++ b/test/enzyme/permute.jl @@ -24,13 +24,11 @@ fTβs = is_ci ? (Duplicated,) : (Duplicated, Const) β = randn(T) p = randindextuple(numind(A)) C = randn!(permute(A, p)) - @static if VERSION ≥ v"1.11.0-rc" - for Tα in rTαs, Tβ in rTβs - EnzymeTestUtils.test_reverse(TensorKit.permute!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "permute! reverse Tα: $Tα, Tβ: $Tβ") - end - for Tα in fTαs, Tβ in fTβs - EnzymeTestUtils.test_forward(TensorKit.permute!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "permute! forward Tα: $Tα, Tβ: $Tβ") - end + for Tα in rTαs, Tβ in rTβs + EnzymeTestUtils.test_reverse(TensorKit.permute!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "permute! reverse Tα: $Tα, Tβ: $Tβ") + end + for Tα in fTαs, Tβ in fTβs + EnzymeTestUtils.test_forward(TensorKit.permute!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "permute! forward Tα: $Tα, Tβ: $Tβ") end end end diff --git a/test/enzyme/tr.jl b/test/enzyme/tr.jl index f42e02ee1..1ba0c2df7 100644 --- a/test/enzyme/tr.jl +++ b/test/enzyme/tr.jl @@ -22,18 +22,12 @@ TDs = is_ci ? (Duplicated,) : (Const, Duplicated) @testset "tr reverse: RT $RT, TD $TD" for RT in rRTs, TD in TDs EnzymeTestUtils.test_reverse(tr, RT, (D1, TD); atol, rtol) EnzymeTestUtils.test_reverse(tr, RT, (D2, TD); atol, rtol) - # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 - @static if VERSION ≥ v"1.11.0-rc" - EnzymeTestUtils.test_reverse(tr, RT, (D3, TD); atol, rtol) - end + EnzymeTestUtils.test_reverse(tr, RT, (D3, TD); atol, rtol) end @testset "tr forward: RT $RT, TD $TD" for RT in fRTs, TD in TDs EnzymeTestUtils.test_forward(tr, RT, (D1, TD); atol, rtol) EnzymeTestUtils.test_forward(tr, RT, (D2, TD); atol, rtol) - # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 - @static if VERSION ≥ v"1.11.0-rc" - EnzymeTestUtils.test_forward(tr, RT, (D3, TD); atol, rtol) - end + EnzymeTestUtils.test_forward(tr, RT, (D3, TD); atol, rtol) end end end diff --git a/test/runtests.jl b/test/runtests.jl index ae8446f2a..de3aed430 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -41,13 +41,14 @@ args = parse_args(ARGS; custom = ["fast"]) # --fast: skip AD tests and inject fast_tests=true into each worker sandbox fast = !isnothing(args.custom["fast"]) -# Enzyme workers end up needing ~4 GB, but the default job count assumes ~2 GB per worker. -# Cap the number of jobs by increasing the per-worker memory budget unless --jobs was given. +# Enzyme workers peak at ~4.5 GB, but the default job count +# assumes ~2 GB per worker. Cap the number of jobs by increasing the per-worker memory budget +# unless --jobs was given; with 4 GB, 16 GB CI runners would run out of memory and be killed. selected = copy(testsuite) ParallelTestRunner.filter_tests!(selected, args) has_enzyme = any(startswith("enzyme"), keys(selected)) if isnothing(args.jobs) && has_enzyme - njobs = clamp(Int(Sys.free_memory() ÷ (4 * Int64(2)^30)), 1, Sys.CPU_THREADS) + njobs = clamp(Int(Sys.free_memory() ÷ (5 * Int64(2)^30)), 1, Sys.CPU_THREADS) # On Windows with Julia 1.10, memory pressure makes recycled workers too slow to start: # Malt gives up if a new worker does not connect within 30 s, which aborts the whole run. Sys.iswindows() && VERSION < v"1.11" && (njobs = min(njobs, 2))