diff --git a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl index 98bc70c40..cadae496a 100644 --- a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl +++ b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl @@ -11,21 +11,6 @@ using MatrixAlgebraKit using TupleTools using Random: AbstractRNG -""" - is_inactive(config, x::Annotation) -> Bool - -Check whether the annotated argument `x` should be treated as inactive, i.e. whether no tangent should be read from or accumulated into `x.dval`. -Under runtime activity, Enzyme may pass an argument that is inactive at run time as a `Duplicated` whose shadow aliases the primal (`x.dval === x.val`); -writing into that shadow would then corrupt the primal. -See https://enzymead.github.io/Enzyme.jl/dev/faq/#faq-runtime-activity. - -!!! note - This is a stopgap until an equivalent helper is available in `EnzymeRules` - (see https://github.com/EnzymeAD/Enzyme.jl/pull/3597#discussion_r4014313289). -""" -@inline is_inactive(config, ::Const) = true -@inline is_inactive(config, x::Annotation) = EnzymeRules.runtime_activity(config) && x.dval === x.val - include("utility.jl") include("linalg.jl") include("indexmanipulations.jl") diff --git a/ext/TensorKitEnzymeExt/tensoroperations.jl b/ext/TensorKitEnzymeExt/tensoroperations.jl index 6b80fd26e..8d6652ba7 100644 --- a/ext/TensorKitEnzymeExt/tensoroperations.jl +++ b/ext/TensorKitEnzymeExt/tensoroperations.jl @@ -22,10 +22,10 @@ function EnzymeRules.augmented_primal( backend::Const, allocator::Const ) where {RT} - Ccache = isa(β, Const) || is_inactive(config, C) ? nothing : copy(C.val) - A_needs_cache = EnzymeRules.overwritten(config)[3] && !(typeof(B) <: Const) && !(typeof(C) <: Const) && !is_inactive(config, A) + Ccache = isa(β, Const) ? nothing : copy(C.val) + A_needs_cache = EnzymeRules.overwritten(config)[3] && !(typeof(B) <: Const) && !(typeof(C) <: Const) Acache = A_needs_cache ? copy(A.val) : nothing - B_needs_cache = EnzymeRules.overwritten(config)[5] && !(typeof(A) <: Const) && !(typeof(C) <: Const) && !is_inactive(config, B) + B_needs_cache = EnzymeRules.overwritten(config)[5] && !(typeof(A) <: Const) && !(typeof(C) <: Const) Bcache = B_needs_cache ? copy(B.val) : nothing AB = if !isa(α, Const) AB = TO.tensorcontract(A.val, pA.val, false, B.val, pB.val, false, pAB.val, One(), backend.val, allocator.val) @@ -65,17 +65,17 @@ function EnzymeRules.reverse( Δα = pullback_dα(α, C, AB) Δβ = pullback_dβ(β, C, Cval) - if !is_inactive(A) && !is_inactive(C) + if !isa(A, Const) TensorKit.blas_contract_pullback_ΔA!( A.dval, C.dval, Aval, pA.val, Bval, pB.val, pAB.val, α.val, backend.val, allocator.val ) # this typically returns nothing end - if !is_inactive(B) && !is_inactive(C) + if !isa(B, Const) TensorKit.blas_contract_pullback_ΔB!( B.dval, C.dval, Aval, pA.val, Bval, pB.val, pAB.val, α.val, backend.val, allocator.val ) # this typically returns nothing end - !is_inactive(C) && pullback_dC!(C.dval, β.val) # this typically returns nothing + !isa(C, Const) && pullback_dC!(C.dval, β.val) # this typically returns nothing return nothing, nothing, nothing, nothing, nothing, nothing, Δα, Δβ, nothing, nothing end