From 364afb71d6e09d13ce6e1762779d13578dc5e0f3 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 22 May 2026 11:43:47 +0200 Subject: [PATCH 1/4] Add Enzyme forward and reverse rules and tests for VectorInterface --- Project.toml | 10 +++---- test/enzyme-vectorinterface/add.jl | 45 ++++++++++++++++++++++++++++ test/enzyme-vectorinterface/inner.jl | 25 ++++++++++++++++ test/enzyme-vectorinterface/scale.jl | 40 +++++++++++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 116 insertions(+), 5 deletions(-) create mode 100644 test/enzyme-vectorinterface/add.jl create mode 100644 test/enzyme-vectorinterface/inner.jl create mode 100644 test/enzyme-vectorinterface/scale.jl diff --git a/Project.toml b/Project.toml index 5342d987f..10d69842f 100644 --- a/Project.toml +++ b/Project.toml @@ -3,6 +3,9 @@ uuid = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec" version = "0.17.0" authors = ["Jutho Haegeman, Lukas Devos"] +[workspace] +projects = ["test", "docs"] + [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" @@ -39,12 +42,9 @@ TensorKitFiniteDifferencesExt = "FiniteDifferences" TensorKitGPUArraysExt = "GPUArrays" TensorKitMooncakeExt = "Mooncake" -[workspace] -projects = ["test", "docs"] - [compat] -Adapt = "4" AMDGPU = "2" +Adapt = "4" CUDA = "6" ChainRulesCore = "1" Dictionaries = "0.4" @@ -62,7 +62,7 @@ Random = "1" ScopedValues = "1.3.0" Strided = "2.6.1" TensorKitSectors = "0.3.7" -TensorOperations = "5.5.2" +TensorOperations = "5.5.2, 5.6" TupleTools = "1.5" VectorInterface = "0.6" julia = "1.10" diff --git a/test/enzyme-vectorinterface/add.jl b/test/enzyme-vectorinterface/add.jl new file mode 100644 index 000000000..ca2413f1d --- /dev/null +++ b/test/enzyme-vectorinterface/add.jl @@ -0,0 +1,45 @@ +using Test, TestExtras +using TensorKit, Enzyme, EnzymeTestUtils +using TensorOperations +using Random + +#spacelist = ad_spacelist(fast_tests) +spacelist = [ad_spacelist(fast_tests)[1]] +eltypes = (Float64, ComplexF64) + +@testset "Enzyme - VectorInterface (add!) $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes + atol = default_tol(T) + rtol = default_tol(T) + + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + α = randn(T) + β = randn(T) + + for TC in (Duplicated,), TA in (Duplicated,) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_reverse(add!, TC, (C, TC), (A, TA); atol, rtol, testset_name = "add! reverse TC $TC TA $TA no α no β") + EnzymeTestUtils.test_forward(add!, TC, (C, TC), (A, TA); atol, rtol, testset_name = "add! forward TC $TC TA $TA no α no β") + for Tα in (Active, Const) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_reverse(add!, TC, (C, TC), (A, TA), (α, Tα); atol, rtol, testset_name = "add! reverse TC $TC TA $TA Tα $Tα no β") + for Tβ in (Active, Const) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_reverse(add!, TC, (C, TC), (A, TA), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add! reverse TC $TC TA $TA Tα $Tα Tβ $Tβ") + end + end + for Tα in (Duplicated, Const) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_forward(add!, TC, (C, TC), (A, TA), (α, Tα); atol, rtol, testset_name = "add! forward TC $TC TA $TA Tα $Tα no β") + for Tβ in (Duplicated, Const) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_forward(add!, TC, (C, TC), (A, TA), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add! forward TC $TC TA $TA Tα $Tα Tβ $Tβ") + end + end + end +end diff --git a/test/enzyme-vectorinterface/inner.jl b/test/enzyme-vectorinterface/inner.jl new file mode 100644 index 000000000..1db51a14b --- /dev/null +++ b/test/enzyme-vectorinterface/inner.jl @@ -0,0 +1,25 @@ +using Test, TestExtras +using TensorKit +using TensorOperations +using Enzyme, EnzymeTestUtils +using Random, FiniteDifferences + +spacelist = ad_spacelist(fast_tests) +eltypes = (Float64, ComplexF64) + +@testset "Enzyme - VectorInterface" begin + @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes + @testset for TC in (Duplicated,), TA in (Duplicated,), f in (identity, adjoint) + atol = default_tol(T) + rtol = default_tol(T) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + for RT in (Active, Const) + EnzymeTestUtils.test_reverse(inner, RT, (f(C), TC), (f(A), TA); atol, rtol) + end + for RT in (Duplicated, Const) + EnzymeTestUtils.test_forward(inner, RT, (f(C), TC), (f(A), TA); atol, rtol) + end + end + end +end diff --git a/test/enzyme-vectorinterface/scale.jl b/test/enzyme-vectorinterface/scale.jl new file mode 100644 index 000000000..622cca252 --- /dev/null +++ b/test/enzyme-vectorinterface/scale.jl @@ -0,0 +1,40 @@ +using Test, TestExtras +using TensorKit +using TensorOperations +using Enzyme, EnzymeTestUtils +using Random + +spacelist = ad_spacelist(fast_tests) +eltypes = (Float64, ComplexF64) + +@testset "Enzyme - VectorInterface (scale!)" begin + @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes + atol = default_tol(T) + rtol = default_tol(T) + α = randn(T) + @testset for TC in (Duplicated,) + for Tα in (Active, Const) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_reverse(scale!, TC, (C, TC), (α, Tα); atol, rtol) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_reverse(scale!, TC, (C', TC), (α, Tα); atol, rtol) + @testset for TA in (Duplicated,), (fc, fa) in ((identity, identity), (adjoint, adjoint)) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_reverse(scale!, TC, (fc(C), TC), (fa(A), TA), (α, Tα); atol, rtol) + end + end + for Tα in (Duplicated, Const) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_forward(scale!, TC, (C, TC), (α, Tα); atol, rtol) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_forward(scale!, TC, (C', TC), (α, Tα); atol, rtol) + @testset for TA in (Duplicated,), (fc, fa) in ((identity, identity), (adjoint, adjoint)) + C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + EnzymeTestUtils.test_forward(scale!, TC, (fc(C), TC), (fa(A), TA), (α, Tα); atol, rtol) + end + end + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 881d538a8..7e76a60e7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,6 +23,7 @@ end if (Sys.isapple() && get(ENV, "CI", "false") == "true") || !isempty(VERSION.prerelease) filter!(!startswith("chainrules") ∘ first, testsuite) filter!(!startswith("mooncake") ∘ first, testsuite) + filter!(!startswith("enzyme") ∘ first, testsuite) end args = parse_args(ARGS; custom = ["fast"]) From 007bba72fc1f45021a8e2ec639c9f0c7cc654eef Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 25 Jun 2026 16:44:15 -0400 Subject: [PATCH 2/4] Handle hashing issue on 1.10 --- test/enzyme-vectorinterface/add.jl | 30 +++++++++++++++++----------- test/enzyme-vectorinterface/inner.jl | 10 ++++++++-- test/enzyme-vectorinterface/scale.jl | 22 ++++++++++++-------- 3 files changed, 40 insertions(+), 22 deletions(-) diff --git a/test/enzyme-vectorinterface/add.jl b/test/enzyme-vectorinterface/add.jl index ca2413f1d..1fa3b0fca 100644 --- a/test/enzyme-vectorinterface/add.jl +++ b/test/enzyme-vectorinterface/add.jl @@ -11,33 +11,39 @@ eltypes = (Float64, ComplexF64) atol = default_tol(T) rtol = default_tol(T) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) α = randn(T) β = randn(T) + # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 + @static if VERSION < v"1.11.0-rc" + CV = V[1] ⊗ V[2] ← V[4] ⊗ V[5] + else + CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5] + end + C = randn(T, CV) + A = randn(T, CV) for TC in (Duplicated,), TA in (Duplicated,) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) + A = randn(T, CV) EnzymeTestUtils.test_reverse(add!, TC, (C, TC), (A, TA); atol, rtol, testset_name = "add! reverse TC $TC TA $TA no α no β") EnzymeTestUtils.test_forward(add!, TC, (C, TC), (A, TA); atol, rtol, testset_name = "add! forward TC $TC TA $TA no α no β") for Tα in (Active, Const) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) + A = randn(T, CV) EnzymeTestUtils.test_reverse(add!, TC, (C, TC), (A, TA), (α, Tα); atol, rtol, testset_name = "add! reverse TC $TC TA $TA Tα $Tα no β") for Tβ in (Active, Const) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) + A = randn(T, CV) EnzymeTestUtils.test_reverse(add!, TC, (C, TC), (A, TA), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add! reverse TC $TC TA $TA Tα $Tα Tβ $Tβ") end end for Tα in (Duplicated, Const) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) + A = randn(T, CV) EnzymeTestUtils.test_forward(add!, TC, (C, TC), (A, TA), (α, Tα); atol, rtol, testset_name = "add! forward TC $TC TA $TA Tα $Tα no β") for Tβ in (Duplicated, Const) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) + A = randn(T, CV) EnzymeTestUtils.test_forward(add!, TC, (C, TC), (A, TA), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add! forward TC $TC TA $TA Tα $Tα Tβ $Tβ") end end diff --git a/test/enzyme-vectorinterface/inner.jl b/test/enzyme-vectorinterface/inner.jl index 1db51a14b..9136fc30b 100644 --- a/test/enzyme-vectorinterface/inner.jl +++ b/test/enzyme-vectorinterface/inner.jl @@ -12,8 +12,14 @@ eltypes = (Float64, ComplexF64) @testset for TC in (Duplicated,), TA in (Duplicated,), f in (identity, adjoint) atol = default_tol(T) rtol = default_tol(T) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 + @static if VERSION < v"1.11.0-rc" + CV = V[1] ⊗ V[2] ← V[4] ⊗ V[5] + else + CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5] + end + C = randn(T, CV) + A = randn(T, CV) for RT in (Active, Const) EnzymeTestUtils.test_reverse(inner, RT, (f(C), TC), (f(A), TA); atol, rtol) end diff --git a/test/enzyme-vectorinterface/scale.jl b/test/enzyme-vectorinterface/scale.jl index 622cca252..cccca12e7 100644 --- a/test/enzyme-vectorinterface/scale.jl +++ b/test/enzyme-vectorinterface/scale.jl @@ -12,26 +12,32 @@ eltypes = (Float64, ComplexF64) atol = default_tol(T) rtol = default_tol(T) α = randn(T) + # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 + @static if VERSION < v"1.11.0-rc" + CV = V[1] ⊗ V[2] ← V[4] ⊗ V[5] + else + CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5] + end @testset for TC in (Duplicated,) for Tα in (Active, Const) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) EnzymeTestUtils.test_reverse(scale!, TC, (C, TC), (α, Tα); atol, rtol) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) EnzymeTestUtils.test_reverse(scale!, TC, (C', TC), (α, Tα); atol, rtol) @testset for TA in (Duplicated,), (fc, fa) in ((identity, identity), (adjoint, adjoint)) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) + A = randn(T, CV) EnzymeTestUtils.test_reverse(scale!, TC, (fc(C), TC), (fa(A), TA), (α, Tα); atol, rtol) end end for Tα in (Duplicated, Const) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) EnzymeTestUtils.test_forward(scale!, TC, (C, TC), (α, Tα); atol, rtol) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) EnzymeTestUtils.test_forward(scale!, TC, (C', TC), (α, Tα); atol, rtol) @testset for TA in (Duplicated,), (fc, fa) in ((identity, identity), (adjoint, adjoint)) - C = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) - A = randn(T, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]) + C = randn(T, CV) + A = randn(T, CV) EnzymeTestUtils.test_forward(scale!, TC, (fc(C), TC), (fa(A), TA), (α, Tα); atol, rtol) end end From 75615ef6396e4ef818ee16539f58564361e60d9f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 25 Jun 2026 19:18:21 -0400 Subject: [PATCH 3/4] Fix norm test for hashing issue on 1.10 --- test/enzyme-linalg/norm.jl | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/test/enzyme-linalg/norm.jl b/test/enzyme-linalg/norm.jl index d12dccc1a..504e51884 100644 --- a/test/enzyme-linalg/norm.jl +++ b/test/enzyme-linalg/norm.jl @@ -14,7 +14,12 @@ 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) - C = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') + # 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 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) From 4cecdbde88dd8d299bbb213d639e03cf01be0f1b Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 26 Jun 2026 00:40:01 -0400 Subject: [PATCH 4/4] Workaround for bad fusion channel --- test/enzyme-vectorinterface/add.jl | 2 +- test/enzyme-vectorinterface/inner.jl | 2 +- test/enzyme-vectorinterface/scale.jl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/enzyme-vectorinterface/add.jl b/test/enzyme-vectorinterface/add.jl index 1fa3b0fca..5d44b2830 100644 --- a/test/enzyme-vectorinterface/add.jl +++ b/test/enzyme-vectorinterface/add.jl @@ -15,7 +15,7 @@ eltypes = (Float64, ComplexF64) β = randn(T) # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 - @static if VERSION < v"1.11.0-rc" + if VERSION < v"1.11.0-rc" && sectortype(eltype(V)) == Trivial CV = V[1] ⊗ V[2] ← V[4] ⊗ V[5] else CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5] diff --git a/test/enzyme-vectorinterface/inner.jl b/test/enzyme-vectorinterface/inner.jl index 9136fc30b..5df4a8017 100644 --- a/test/enzyme-vectorinterface/inner.jl +++ b/test/enzyme-vectorinterface/inner.jl @@ -13,7 +13,7 @@ eltypes = (Float64, ComplexF64) 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" + if VERSION < v"1.11.0-rc" && sectortype(eltype(V)) == Trivial CV = V[1] ⊗ V[2] ← V[4] ⊗ V[5] else CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5] diff --git a/test/enzyme-vectorinterface/scale.jl b/test/enzyme-vectorinterface/scale.jl index cccca12e7..8d3f109d5 100644 --- a/test/enzyme-vectorinterface/scale.jl +++ b/test/enzyme-vectorinterface/scale.jl @@ -13,7 +13,7 @@ eltypes = (Float64, ComplexF64) rtol = default_tol(T) α = randn(T) # see https://github.com/QuantumKitHub/TensorKit.jl/issues/457 - @static if VERSION < v"1.11.0-rc" + if VERSION < v"1.11.0-rc" && sectortype(eltype(V)) == Trivial CV = V[1] ⊗ V[2] ← V[4] ⊗ V[5] else CV = V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5]