Reverse/forward Enzyme rules, refactoring, and tests for planar ops

Project.toml

@@ -48,7 +48,7 @@ Adapt = "4"
 CUDA = "6"
 ChainRulesCore = "1"
 Dictionaries = "0.4"
-Enzyme = "0.13.157"
+Enzyme = "0.13.195"
 EnzymeTestUtils = "0.2.8"
 FiniteDifferences = "0.12"
 GPUArrays = "11.4.1"

ext/TensorKitEnzymeExt/linalg.jl

@@ -11,9 +11,9 @@ function EnzymeRules.augmented_primal(
         α::Annotation,
         β::Annotation,
     ) where {RT}
-    cacheC = !isa(β, Const) && copy(C.val)
-    cacheA = !isa(B, Const) && EnzymeRules.overwritten(config)[3] ? copy(A.val) : nothing
-    cacheB = !isa(A, Const) && EnzymeRules.overwritten(config)[4] ? copy(B.val) : nothing
+    cacheC = !isa(β, Const) ? copy(C.val) : C.val
+    cacheA = !isa(B, Const) && EnzymeRules.overwritten(config)[3] ? copy(A.val) : A.val
+    cacheB = !isa(A, Const) && EnzymeRules.overwritten(config)[4] ? copy(B.val) : B.val
     AB = if !isa(α, Const)
         AB = A.val * B.val
         add!(C.val, AB, α.val, β.val)
@@ -39,16 +39,7 @@ function EnzymeRules.reverse(
         α::Annotation{<:Number},
         β::Annotation{<:Number},
     ) where {RT}
-    if RT <: Const
-        Δα = isa(α, Const) ? nothing : zero(α.val)
-        Δβ = isa(β, Const) ? nothing : zero(β.val)
-        return (nothing, nothing, nothing, Δα, Δβ)
-    end
-    cacheC, cacheA, cacheB, AB = cache
-    Cval = something(cacheC, C.val)
-    Aval = something(cacheA, A.val)
-    Bval = something(cacheB, B.val)
-
+    Cval, Aval, Bval, AB = cache
     !isa(A, Const) && !isa(C, Const) && TK.project_mul!(A.dval, C.dval, Bval', conj(α.val), One())
     !isa(B, Const) && !isa(C, Const) && TK.project_mul!(B.dval, Aval', C.dval, conj(α.val), One())
     Δαr = pullback_dα(α, C, AB)

ext/TensorKitEnzymeExt/utility.jl

@@ -15,9 +15,17 @@ pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β))
 # Ignore derivatives
 # ------------------
 
+@inline EnzymeRules.inactive_type(::Type{TensorKit.GenericFusion}) = true
+@inline EnzymeRules.inactive_type(::Type{TensorKit.UniqueFusion}) = true
+@inline EnzymeRules.inactive_type(::Type{TensorKit.MultipleFusion}) = true
+@inline EnzymeRules.inactive_type(::Type{TensorKit.SimpleFusion}) = true
+@inline EnzymeRules.inactive_type(::Type{TensorKit.MultiplicityFreeFusion}) = true
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.FusionTree}) = true
+@inline EnzymeRules.inactive_type(::Type{<:TensorKit.FusionTreeDict}) = true
+@inline EnzymeRules.inactive_type(::Type{<:TensorKit.FusionTreeBlock}) = true
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true
 @inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true
+@inline EnzymeRules.inactive_type(::Type{<:TensorKit.LRU}) = true
 
 function EnzymeRules.augmented_primal(
         config::EnzymeRules.RevConfigWidth{1},
@@ -34,7 +42,7 @@ function EnzymeRules.augmented_primal(
     else
         nothing
     end
-    return EnzymeRules.AugmentedReturn(ret, dret, dret)
+    return EnzymeRules.AugmentedReturn(ret, dret, nothing)
 end
 
 function EnzymeRules.reverse(
@@ -129,6 +137,10 @@ function EnzymeRules.forward(
     end
 end
 
+@inline EnzymeRules.inactive(::typeof(TensorKit.planar_trace), ::TensorKit.FusionTreePair, ::TensorKit.Index2Tuple, ::TensorKit.Index2Tuple) = nothing
+@inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any) = nothing
+@inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any, ::Any) = nothing
+@inline EnzymeRules.inactive(::typeof(TensorKit.fusiontrees), ::Any, ::Any, ::Any) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any, ::Any) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.artin_braid), ::Any, ::Any) = nothing

ext/TensorKitMooncakeExt/TensorKitMooncakeExt.jl

@@ -19,7 +19,6 @@ include("linalg.jl")
 include("indexmanipulations.jl")
 include("vectorinterface.jl")
 include("tensoroperations.jl")
-include("planaroperations.jl")
 include("factorizations.jl")
 
 end

ext/TensorKitMooncakeExt/utility.jl

@@ -20,6 +20,7 @@ Mooncake.tangent_type(::Type{<:HomSpace}) = Mooncake.NoTangent
 
 @zero_derivative DefaultCtx Tuple{typeof(TensorOperations.tensorcontract_structure), AbstractTensorMap, Index2Tuple, Bool, AbstractTensorMap, Index2Tuple, Bool, Index2Tuple}
 
+@zero_derivative DefaultCtx Tuple{typeof(TensorKit.planar_trace), TensorKit.FusionTreePair, Index2Tuple, Index2Tuple}
 @zero_derivative DefaultCtx Tuple{typeof(TensorKit.has_shared_permute), AbstractTensorMap, Index2Tuple}
 @zero_derivative DefaultCtx Tuple{typeof(TensorKit.select), HomSpace, Index2Tuple}
 @zero_derivative DefaultCtx Tuple{typeof(TensorKit.flip), HomSpace, Any}

src/TensorKit.jl

@@ -285,7 +285,6 @@ include("planar/planaroperations.jl")
 # ------------------------
 include("auxiliary/ad.jl")
 include("pullbacks/tensoroperations.jl")
-
 include("pullbacks/indexmanipulations.jl")
 
 end

test/Project.toml

@@ -43,3 +43,4 @@ ParallelTestRunner = "2"
 Test = "1"
 TestExtras = "0.2,0.3"
 Zygote = "0.7"
+

test/enzyme-planaroperations/planarcontract.jl

@@ -0,0 +1,78 @@
+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

test/enzyme-planaroperations/planartrace.jl

@@ -0,0 +1,41 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using VectorInterface: Zero, One
+using Enzyme, EnzymeTestUtils
+using Random
+
+spacelist = ad_spacelist(fast_tests)
+#eltypes = (Float64, ComplexF64)
+eltypes = (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