From 657ab82dc7a0f1439c9f5268c7357757bda9817a Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 30 Jul 2026 12:19:12 +0200 Subject: [PATCH 1/4] Try combining the mul/planarcontract tests --- test/enzyme-linalg/mul.jl | 110 +++++++++++------- .../enzyme-planaroperations/planarcontract.jl | 78 ------------- 2 files changed, 65 insertions(+), 123 deletions(-) delete mode 100644 test/enzyme-planaroperations/planarcontract.jl diff --git a/test/enzyme-linalg/mul.jl b/test/enzyme-linalg/mul.jl index 1250df342..05da0aab9 100644 --- a/test/enzyme-linalg/mul.jl +++ b/test/enzyme-linalg/mul.jl @@ -9,61 +9,81 @@ eltypes = (Float64, ComplexF64) is_ci = get(ENV, "CI", "false") == "true" -@timedtestset verbose = true "Enzyme - LinearAlgebra (mul):" begin +rTs(::Number, is_ci::Bool) = is_ci ? (Active,) : (Active, Const) +rTs(::Zero, is_ci::Bool) = (Const,) +rTs(::One, is_ci::Bool) = (Const,) +fTs(::Number, is_ci::Bool) = is_ci ? (Duplicated,) : (Duplicated, Const) +fTs(::Zero, is_ci::Bool) = (Const,) +fTs(::One, is_ci::Bool) = (Const,) + +@timedtestset verbose = true "Enzyme - LinearAlgebra (mul) and planarcontract!:" begin @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes 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 - A = randn(T, codomain(C) ← V[5]' ⊗ V[4]') - B = randn(T, domain(A) ← domain(C)) zero_αβs = ((Zero(), Zero()), (randn(T), Zero()), (Zero(), randn(T))) αβs = !is_ci ? vcat(zero_αβs..., (randn(T), randn(T))) : ((randn(T), randn(T)),) - for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,) - for (α, β) in αβs - rTαs = if α === Zero() - (Const,) - elseif !is_ci - (Active, Const) - else - (Active,) + @timedtestset "mul" begin + C = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') + A = randn(T, codomain(C) ← V[5]' ⊗ V[4]') + B = randn(T, domain(A) ← domain(C)) + for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,) + for (α, β) in αβs + rTαs = rTs(α) + rTβs = rTs(β) + fTαs = fTs(α) + fTβs = fTs(β) + for Tα in rTαs, Tβ in rTβs + EnzymeTestUtils.test_reverse(mul!, TC, (C, TC), (A, TA), (B, TB), (α, Tα), (β, Tβ); atol, rtol, testset_name = "mul! reverse Tα $Tα, Tβ $Tβ") + end + for Tα in fTαs, Tβ in fTβs + EnzymeTestUtils.test_forward(mul!, TC, (C, TC), (A, TA), (B, TB), (α, Tα), (β, Tβ); atol, rtol, testset_name = "mul! forward Tα $Tα, Tβ $Tβ") + end end - rTβs = if β === Zero() - (Const,) - elseif !is_ci - (Active, Const) - else - (Active,) + if !is_ci + EnzymeTestUtils.test_reverse(mul!, TC, (C, TC), (A, TA), (B, TB); atol, rtol, testset_name = "mul! reverse no α no β") + EnzymeTestUtils.test_forward(mul!, TC, (C, TC), (A, TA), (B, TB); atol, rtol, testset_name = "mul! forward no α no β") end + end + end + @timedtestset "planarcontract!" begin + V1, V2, V3, V4, V5 = V + k1 = 3 + k2 = 2 + k3 = 3 + k′ = rand(0:(k1 + k2)) + pA = randcircshift(k′, k1 + k2 - k′, k1) + ipA = _repartition(invperm(linearize(pA)), k′) + k′ = rand(0:(k2 + k3)) + pB = randcircshift(k′, k2 + k3 - k′, k2) + ipB = _repartition(invperm(linearize(pB)), k′) + # TODO: primal value already is broken for this? + # pAB = randcircshift(k1, k3) + pAB = _repartition(tuple((1:(k1 + k3))...), k1) + + α_ = randn(T) + β_ = randn(T) + + A = randn(T, permute(V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)', ipA)) + B = randn(T, permute((V4 ⊗ V5)' ← V1 ⊗ V2 ⊗ V3, ipB)) + C = randn!( + TensorOperations.tensoralloc_contract( + T, A, pA, false, B, pB, false, pAB, Val(false) + ) + ) + αβs = !is_ci ? ((One(), Zero()), (α_, β_)) : ((α_, β_),) + @testset for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,), (α, β) in αβs + rTαs = rTs(α) + rTβs = rTs(β) + fTαs = fTs(α) + fTβs = fTs(β) for Tα in rTαs, Tβ in rTβs - EnzymeTestUtils.test_reverse(mul!, TC, (C, TC), (A, TA), (B, TB), (α, Tα), (β, Tβ); atol, rtol, testset_name = "mul! reverse Tα $Tα, Tβ $Tβ") - end - fTαs = if α === Zero() - (Const,) - elseif !is_ci - (Duplicated, Const) - else - (Duplicated,) + EnzymeTestUtils.test_reverse(TensorKit.planarcontract!, TC, (C, TC), (A, TA), (pA, Const), (B, TB), (pB, Const), (pAB, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "planarcontract! Tα = $Tα, Tβ = $Tβ") end - fTβs = if β === Zero() - (Const,) - elseif !is_ci - (Duplicated, Const) - else - (Duplicated,) - end - for Tα in fTαs, Tβ in fTβs - EnzymeTestUtils.test_forward(mul!, TC, (C, TC), (A, TA), (B, TB), (α, Tα), (β, Tβ); atol, rtol, testset_name = "mul! forward Tα $Tα, Tβ $Tβ") - end - end - if !is_ci - EnzymeTestUtils.test_reverse(mul!, TC, (C, TC), (A, TA), (B, TB); atol, rtol, testset_name = "mul! reverse no α no β") - EnzymeTestUtils.test_forward(mul!, TC, (C, TC), (A, TA), (B, TB); atol, rtol, testset_name = "mul! forward no α no β") + # TODO broken internally in Enzyme? + #for Tα in fTαs, Tβ in fTβs + # EnzymeTestUtils.test_forward(TensorKit.planarcontract!, TC, (C, TC), (A, TA), (pA, Const), (B, TB), (pB, Const), (pAB, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "planarcontract! Tα = $Tα, Tβ = $Tβ") + #end end end end diff --git a/test/enzyme-planaroperations/planarcontract.jl b/test/enzyme-planaroperations/planarcontract.jl deleted file mode 100644 index 9ccf4e95a..000000000 --- a/test/enzyme-planaroperations/planarcontract.jl +++ /dev/null @@ -1,78 +0,0 @@ -using Test, TestExtras -using TensorKit -using TensorOperations -using VectorInterface: Zero, One -using Enzyme, EnzymeTestUtils -using Random - -is_ci = get(ENV, "CI", "false") == "true" - -spacelist = ad_spacelist(fast_tests) -eltypes = (Float64, ComplexF64) - -@timedtestset "Enzyme - PlanarOperations (planarcontract): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes - atol = default_tol(T) - rtol = default_tol(T) - V1, V2, V3, V4, V5 = V - k1 = 3 - k2 = 2 - k3 = 3 - k′ = rand(0:(k1 + k2)) - pA = randcircshift(k′, k1 + k2 - k′, k1) - ipA = _repartition(invperm(linearize(pA)), k′) - k′ = rand(0:(k2 + k3)) - pB = randcircshift(k′, k2 + k3 - k′, k2) - ipB = _repartition(invperm(linearize(pB)), k′) - # TODO: primal value already is broken for this? - # pAB = randcircshift(k1, k3) - pAB = _repartition(tuple((1:(k1 + k3))...), k1) - - α_ = randn(T) - β_ = randn(T) - - A = randn(T, permute(V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)', ipA)) - B = randn(T, permute((V4 ⊗ V5)' ← V1 ⊗ V2 ⊗ V3, ipB)) - C = randn!( - TensorOperations.tensoralloc_contract( - T, A, pA, false, B, pB, false, pAB, Val(false) - ) - ) - αβs = !is_ci ? ((One(), Zero()), (α_, β_)) : ((α_, β_),) - @testset for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,), (α, β) in αβs - rTαs = if (α === Zero() || α === One()) - (Const,) - elseif !is_ci - (Active, Const) - else - (Active,) - end - rTβs = if (β === Zero() || β === One()) - (Const,) - elseif !is_ci - (Active, Const) - else - (Active,) - end - for Tα in rTαs, Tβ in rTβs - EnzymeTestUtils.test_reverse(TensorKit.planarcontract!, TC, (C, TC), (A, TA), (pA, Const), (B, TB), (pB, Const), (pAB, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "planarcontract! Tα = $Tα, Tβ = $Tβ") - end - fTαs = if (α === Zero() || α === One()) - (Const,) - elseif !is_ci - (Duplicated, Const) - else - (Duplicated,) - end - fTβs = if (β === Zero() || β === One()) - (Const,) - elseif !is_ci - (Duplicated, Const) - else - (Duplicated,) - end - # TODO broken internally in Enzyme? - #for Tα in fTαs, Tβ in fTβs - # EnzymeTestUtils.test_forward(TensorKit.planarcontract!, TC, (C, TC), (A, TA), (pA, Const), (B, TB), (pB, Const), (pAB, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "planarcontract! Tα = $Tα, Tβ = $Tβ") - #end - end -end From cae775f80a926ce50e356c30e2df0a45194f7e55 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 30 Jul 2026 14:06:41 +0200 Subject: [PATCH 2/4] Include TO --- test/enzyme-linalg/mul.jl | 1 + test/enzyme-planaroperations/planartrace.jl | 3 +-- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/enzyme-linalg/mul.jl b/test/enzyme-linalg/mul.jl index 05da0aab9..4492707b7 100644 --- a/test/enzyme-linalg/mul.jl +++ b/test/enzyme-linalg/mul.jl @@ -1,5 +1,6 @@ using Test, TestExtras using TensorKit +using TensorOperations using VectorInterface using Enzyme, EnzymeTestUtils using Random diff --git a/test/enzyme-planaroperations/planartrace.jl b/test/enzyme-planaroperations/planartrace.jl index 193c0b9f0..7859107e4 100644 --- a/test/enzyme-planaroperations/planartrace.jl +++ b/test/enzyme-planaroperations/planartrace.jl @@ -6,8 +6,7 @@ using Enzyme, EnzymeTestUtils using Random spacelist = ad_spacelist(fast_tests) -#eltypes = (Float64, ComplexF64) -eltypes = (ComplexF64,) +eltypes = (Float64, ComplexF64) @timedtestset "Enzyme - PlanarOperations (planartrace): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes atol = default_tol(T) From f51ef8c44c75b0a26e7a65bb386669545e3ad91d Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 30 Jul 2026 17:45:51 +0200 Subject: [PATCH 3/4] Dumb typo --- test/enzyme-linalg/mul.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/test/enzyme-linalg/mul.jl b/test/enzyme-linalg/mul.jl index 4492707b7..ad8b1fcd3 100644 --- a/test/enzyme-linalg/mul.jl +++ b/test/enzyme-linalg/mul.jl @@ -30,10 +30,10 @@ fTs(::One, is_ci::Bool) = (Const,) B = randn(T, domain(A) ← domain(C)) for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,) for (α, β) in αβs - rTαs = rTs(α) - rTβs = rTs(β) - fTαs = fTs(α) - fTβs = fTs(β) + rTαs = rTs(α, is_ci) + rTβs = rTs(β, is_ci) + fTαs = fTs(α, is_ci) + fTβs = fTs(β, is_ci) for Tα in rTαs, Tβ in rTβs EnzymeTestUtils.test_reverse(mul!, TC, (C, TC), (A, TA), (B, TB), (α, Tα), (β, Tβ); atol, rtol, testset_name = "mul! reverse Tα $Tα, Tβ $Tβ") end @@ -74,10 +74,10 @@ fTs(::One, is_ci::Bool) = (Const,) ) αβs = !is_ci ? ((One(), Zero()), (α_, β_)) : ((α_, β_),) @testset for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,), (α, β) in αβs - rTαs = rTs(α) - rTβs = rTs(β) - fTαs = fTs(α) - fTβs = fTs(β) + rTαs = rTs(α, is_ci) + rTβs = rTs(β, is_ci) + fTαs = fTs(α, is_ci) + fTβs = fTs(β, is_ci) for Tα in rTαs, Tβ in rTβs EnzymeTestUtils.test_reverse(TensorKit.planarcontract!, TC, (C, TC), (A, TA), (pA, Const), (B, TB), (pB, Const), (pAB, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "planarcontract! Tα = $Tα, Tβ = $Tβ") end From eab9a41cfd2489723042824a4ca06439540f5cc7 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 30 Jul 2026 19:14:42 +0200 Subject: [PATCH 4/4] Combine planartrace with trace_permute --- test/enzyme-planaroperations/planartrace.jl | 40 --------------------- test/enzyme-vi-to/trace.jl | 28 +++++++++++++++ 2 files changed, 28 insertions(+), 40 deletions(-) delete mode 100644 test/enzyme-planaroperations/planartrace.jl diff --git a/test/enzyme-planaroperations/planartrace.jl b/test/enzyme-planaroperations/planartrace.jl deleted file mode 100644 index 7859107e4..000000000 --- a/test/enzyme-planaroperations/planartrace.jl +++ /dev/null @@ -1,40 +0,0 @@ -using Test, TestExtras -using TensorKit -using TensorOperations -using VectorInterface: Zero, One -using Enzyme, EnzymeTestUtils -using Random - -spacelist = ad_spacelist(fast_tests) -eltypes = (Float64, ComplexF64) - -@timedtestset "Enzyme - PlanarOperations (planartrace): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes - atol = default_tol(T) - rtol = default_tol(T) - for _ in 1:5 - k1 = rand(0:2) - k2 = rand(0:1) - V1 = map(v -> rand(Bool) ? v' : v, rand(V, k1)) - V2 = map(v -> rand(Bool) ? v' : v, rand(V, k2)) - V3 = prod(x -> x ⊗ x', V2[1:k2]; init = one(V[1])) - V4 = prod(x -> x ⊗ x', V2[(k2 + 1):end]; init = one(V[1])) - - k′ = rand(0:(k1 + 2k2)) - (_p, _q) = randcircshift(k′, k1 + 2k2 - k′, k1) - p = _repartition(_p, rand(0:k1)) - q = (tuple(_q[1:2:end]...), tuple(_q[2:2:end]...)) - if !all(isempty, p) && !all(isempty, q) - α = randn(T) - β = randn(T) - ip = _repartition(invperm(linearize((_p, _q))), k′) - A = randn(T, permute(prod(V1) ⊗ V3 ← V4, ip)) - C = randn!(TensorOperations.tensoralloc_add(T, A, p, false, Val(false))) - for Tα in (Const, Active), Tβ in (Const, Active) - EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace reverse Tα $Tα Tβ $Tβ") - end - for Tα in (Const, Duplicated), Tβ in (Const, Duplicated) - EnzymeTestUtils.test_forward(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace forward Tα $Tα Tβ $Tβ") - end - end - end -end diff --git a/test/enzyme-vi-to/trace.jl b/test/enzyme-vi-to/trace.jl index 958af2e11..82f1f9053 100644 --- a/test/enzyme-vi-to/trace.jl +++ b/test/enzyme-vi-to/trace.jl @@ -56,5 +56,33 @@ TAs = is_ci ? (Duplicated,) : (Const, Duplicated) end end end + @timedtestset "planartrace!" begin + for _ in 1:5 + k1 = rand(0:2) + k2 = rand(0:1) + V1 = map(v -> rand(Bool) ? v' : v, rand(V, k1)) + V2 = map(v -> rand(Bool) ? v' : v, rand(V, k2)) + V3 = prod(x -> x ⊗ x', V2[1:k2]; init = one(V[1])) + V4 = prod(x -> x ⊗ x', V2[(k2 + 1):end]; init = one(V[1])) + + k′ = rand(0:(k1 + 2k2)) + (_p, _q) = randcircshift(k′, k1 + 2k2 - k′, k1) + p = _repartition(_p, rand(0:k1)) + q = (tuple(_q[1:2:end]...), tuple(_q[2:2:end]...)) + if !all(isempty, p) && !all(isempty, q) + α = randn(T) + β = randn(T) + ip = _repartition(invperm(linearize((_p, _q))), k′) + A = randn(T, permute(prod(V1) ⊗ V3 ← V4, ip)) + C = randn!(TensorOperations.tensoralloc_add(T, A, p, false, Val(false))) + for Tα in (Const, Active), Tβ in (Const, Active) + EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace reverse Tα $Tα Tβ $Tβ") + end + for Tα in (Const, Duplicated), Tβ in (Const, Duplicated) + EnzymeTestUtils.test_forward(TensorKit.planartrace!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Tα), (β, Tβ), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol, testset_name = "planartrace forward Tα $Tα Tβ $Tβ") + end + end + end + end end end