diff --git a/Project.toml b/Project.toml index 4344242d5..5edabec65 100644 --- a/Project.toml +++ b/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" diff --git a/ext/TensorKitEnzymeExt/linalg.jl b/ext/TensorKitEnzymeExt/linalg.jl index 61022ab91..35e3eaaf8 100644 --- a/ext/TensorKitEnzymeExt/linalg.jl +++ b/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) diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index dd99d96e8..c0f178cf0 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/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 diff --git a/ext/TensorKitMooncakeExt/TensorKitMooncakeExt.jl b/ext/TensorKitMooncakeExt/TensorKitMooncakeExt.jl index 7c0492239..a0ca20f68 100644 --- a/ext/TensorKitMooncakeExt/TensorKitMooncakeExt.jl +++ b/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 diff --git a/ext/TensorKitMooncakeExt/planaroperations.jl b/ext/TensorKitMooncakeExt/planaroperations.jl deleted file mode 100644 index 48d21a78f..000000000 --- a/ext/TensorKitMooncakeExt/planaroperations.jl +++ /dev/null @@ -1,92 +0,0 @@ -# planartrace! -# ------------ -# TODO: Fix planartrace pullback -# This implementation is slightly more involved than its non-planar counterpart -# this is because we lack a general `pAB` argument in `planarcontract`, and need -# to keep things planar along the way. -# In particular, we can't simply tensor product with multiple identities in one go -# if they aren't "contiguous", e.g. p = ((1, 4, 5), ()), q = ((2, 6), (3, 7)) - -# @is_primitive( -# DefaultCtx, -# ReverseMode, -# Tuple{ -# typeof(TensorKit.planartrace!), -# AbstractTensorMap, -# AbstractTensorMap, Index2Tuple, Index2Tuple, -# Number, Number, -# Any, Any, -# } -# ) - -# function Mooncake.rrule!!( -# ::CoDual{typeof(TensorKit.planartrace!)}, -# C_ΔC::CoDual{<:AbstractTensorMap}, -# A_ΔA::CoDual{<:AbstractTensorMap}, p_Δp::CoDual{<:Index2Tuple}, q_Δq::CoDual{<:Index2Tuple}, -# α_Δα::CoDual{<:Number}, β_Δβ::CoDual{<:Number}, -# backend_Δbackend::CoDual, allocator_Δallocator::CoDual -# ) -# # prepare arguments -# C, ΔC = arrayify(C_ΔC) -# A, ΔA = arrayify(A_ΔA) -# p = primal(p_Δp) -# q = primal(q_Δq) -# α, β = primal.((α_Δα, β_Δβ)) -# backend, allocator = primal.((backend_Δbackend, allocator_Δallocator)) -# -# # primal call -# C_cache = copy(C) -# TensorKit.planartrace!(C, A, p, q, α, β, backend, allocator) -# -# function planartrace_pullback(::NoRData) -# copy!(C, C_cache) -# -# ΔAr = planartrace_pullback_ΔA!(ΔA, ΔC, A, p, q, α, backend, allocator) # this typically returns NoRData() -# Δαr = planartrace_pullback_Δα(ΔC, A, p, q, α, backend, allocator) -# Δβr = pullback_dβ(ΔC, C, β) -# ΔCr = pullback_dC!(ΔC, β) # this typically returns NoRData() -# -# return NoRData(), -# ΔCr, ΔAr, NoRData(), NoRData(), -# Δαr, Δβr, NoRData(), NoRData() -# end -# -# return C_ΔC, planartrace_pullback -# end - -# function planartrace_pullback_dA!( -# ΔA, ΔC, A, p, q, α, backend, allocator -# ) -# if length(q[1]) == 0 -# ip = invperm(linearize(p)) -# pΔA = TK._repartition(ip, A) -# TK.transpose!(ΔA, ΔC, pΔA, conj(α), One(), backend, allocator) -# return NoRData() -# end -# # if length(q[1]) == 1 -# # ip = invperm((p[1]..., q[2]..., p[2]..., q[1]...)) -# # pdA = TK._repartition(ip, A) -# # E = one!(TO.tensoralloc_add(scalartype(A), A, q, false)) -# # twist!(E, filter(x -> !isdual(space(E, x)), codomainind(E))) -# # # pE = ((), TK.trivtuple(TO.numind(q))) -# # # pΔC = (TK.trivtuple(TO.numind(p)), ()) -# # TensorKit.planaradd!(ΔA, ΔC ⊗ E, pdA, conj(α), One(), backend, allocator) -# # return NoRData() -# # end -# error("The reverse rule for `planartrace` is not yet implemented") -# end -# -# function planartrace_pullback_dα( -# ΔC, A, p, q, α, backend, allocator -# ) -# Tdα = Mooncake.rdata_type(Mooncake.tangent_type(typeof(α))) -# Tdα === NoRData && return NoRData() -# -# # TODO: this result might be easier to compute as: -# # C′ = βC + α * trace(A) ⟹ At = (C′ - βC) / α -# At = TO.tensoralloc_add(scalartype(A), A, p, false, Val(true), allocator) -# TensorKit.planartrace!(At, A, p, q, One(), Zero(), backend, allocator) -# Δα = project_scalar(α, inner(At, ΔC)) -# TO.tensorfree!(At, allocator) -# return Δα -# end diff --git a/ext/TensorKitMooncakeExt/utility.jl b/ext/TensorKitMooncakeExt/utility.jl index 6d198e901..1bedf566c 100644 --- a/ext/TensorKitMooncakeExt/utility.jl +++ b/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} diff --git a/src/TensorKit.jl b/src/TensorKit.jl index f20859dc6..87a3a2380 100644 --- a/src/TensorKit.jl +++ b/src/TensorKit.jl @@ -285,7 +285,6 @@ include("planar/planaroperations.jl") # ------------------------ include("auxiliary/ad.jl") include("pullbacks/tensoroperations.jl") - include("pullbacks/indexmanipulations.jl") end diff --git a/test/Project.toml b/test/Project.toml index d8acf9e80..ca00312e2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -43,3 +43,4 @@ ParallelTestRunner = "2" Test = "1" TestExtras = "0.2,0.3" Zygote = "0.7" + diff --git a/test/enzyme-planaroperations/planarcontract.jl b/test/enzyme-planaroperations/planarcontract.jl new file mode 100644 index 000000000..9ccf4e95a --- /dev/null +++ b/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 diff --git a/test/enzyme-planaroperations/planartrace.jl b/test/enzyme-planaroperations/planartrace.jl new file mode 100644 index 000000000..193c0b9f0 --- /dev/null +++ b/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