diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index c76f17c70..f772348e5 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -116,33 +116,6 @@ function EnzymeRules.reverse( return nothing, nothing, nothing, nothing, Δαr, Δβr, map(Returns(nothing), ba)... end -function EnzymeRules.augmented_primal( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(twist!)}, - ::Type{RT}, - t::Annotation{<:AbstractTensorMap}, - inds::Const; - inv::Bool = false - ) where {RT} - twist!(t.val, inds.val; inv) - primal = EnzymeRules.needs_primal(config) ? t.val : nothing - shadow = EnzymeRules.needs_shadow(config) ? t.dval : nothing - return EnzymeRules.AugmentedReturn(primal, shadow, nothing) -end - -function EnzymeRules.reverse( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(twist!)}, - ::Type{RT}, - cache, - t::Annotation{<:AbstractTensorMap}, - inds::Const; - inv::Bool = false - ) where {RT} - !isa(t, Const) && twist!(t.dval, inds.val; inv = !inv) - return (nothing, nothing) -end - function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(flip)}, diff --git a/test/enzyme-indexmanipulations-flip-twist/twist.jl b/test/enzyme-indexmanipulations-flip-twist/twist.jl index 6dedcdd5a..594b579be 100644 --- a/test/enzyme-indexmanipulations-flip-twist/twist.jl +++ b/test/enzyme-indexmanipulations-flip-twist/twist.jl @@ -20,6 +20,10 @@ if !Sys.iswindows() && VERSION > v"1.11.0-rc" EnzymeTestUtils.test_reverse(twist!, TA, (copy(A), TA), ([1, 3], Const); atol, rtol, fkwargs = (inv = true,)) EnzymeTestUtils.test_reverse(twist!, TA, (copy(A), TA), (1, Const); atol, rtol) EnzymeTestUtils.test_reverse(twist!, TA, (copy(A), TA), ([1, 3], Const); atol, rtol) + EnzymeTestUtils.test_forward(twist!, TA, (copy(A), TA), (1, Const); atol, rtol, fkwargs = (inv = false,)) + EnzymeTestUtils.test_forward(twist!, TA, (copy(A), TA), ([1, 3], Const); atol, rtol, fkwargs = (inv = true,)) + EnzymeTestUtils.test_forward(twist!, TA, (copy(A), TA), (1, Const); atol, rtol) + EnzymeTestUtils.test_forward(twist!, TA, (copy(A), TA), ([1, 3], Const); atol, rtol) end end end