diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index 0adbb0600..e93c06dac 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -148,3 +148,24 @@ function EnzymeRules.reverse( end return (nothing, nothing) end + +function EnzymeRules.forward( + config::EnzymeRules.FwdConfigWidth{1}, + func::Const{typeof(flip)}, + ::Type{RT}, + t::Annotation{<:AbstractTensorMap}, + inds::Annotation; + inv::Bool = false, + ) where {RT} + t′ = flip(t.val, inds.val; inv) + dt′ = !isa(t, Const) ? flip(t.dval, inds.val; inv) : make_zero(t.val) + if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) + return Duplicated(t′, dt′) + elseif EnzymeRules.needs_primal(config) + return t′ + elseif EnzymeRules.needs_shadow(config) + return dt′ + else + return nothing + end +end diff --git a/test/enzyme-indexmanipulations-flip-twist/flip.jl b/test/enzyme-indexmanipulations-flip-twist/flip.jl index d410673e4..e2de670fc 100644 --- a/test/enzyme-indexmanipulations-flip-twist/flip.jl +++ b/test/enzyme-indexmanipulations-flip-twist/flip.jl @@ -6,19 +6,22 @@ using Random spacelist = ad_spacelist(fast_tests) eltypes = (Float64, ComplexF64) -if !Sys.iswindows() && VERSION > v"1.11.0-rc" - @timedtestset "Enzyme - Index Manipulations (flip):" begin - @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T) TA ($TA)" for V in spacelist, T in eltypes, TA in (Duplicated,) - atol = default_tol(T) - rtol = default_tol(T) - has_braiding = BraidingStyle(sectortype(eltype(V))) isa HasBraiding - if has_braiding - A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') - EnzymeTestUtils.test_reverse(flip, TA, (A, TA), (1, Const); atol, rtol, fkwargs = (inv = false,)) - EnzymeTestUtils.test_reverse(flip, TA, (A, TA), [1, 3]; atol, rtol, fkwargs = (inv = true,)) - EnzymeTestUtils.test_reverse(flip, TA, (A, TA), (1, Const); atol, rtol) - EnzymeTestUtils.test_reverse(flip, TA, (A, TA), ([1, 3], Const); atol, rtol) - end +@timedtestset "Enzyme - Index Manipulations (flip):" begin + @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T) TA ($TA)" for V in spacelist, T in eltypes, TA in (Duplicated,) + atol = default_tol(T) + rtol = default_tol(T) + has_braiding = BraidingStyle(sectortype(eltype(V))) isa HasBraiding + if has_braiding + A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') + EnzymeTestUtils.test_reverse(flip, TA, (A, TA), (1, Const); atol, rtol, fkwargs = (inv = false,)) + EnzymeTestUtils.test_reverse(flip, TA, (A, TA), ([1, 3], Const); atol, rtol, fkwargs = (inv = true,)) + EnzymeTestUtils.test_reverse(flip, TA, (A, TA), (1, Const); atol, rtol) + EnzymeTestUtils.test_reverse(flip, TA, (A, TA), ([1, 3], Const); atol, rtol) + + EnzymeTestUtils.test_forward(flip, TA, (A, TA), (1, Const); atol, rtol, fkwargs = (inv = false,)) + EnzymeTestUtils.test_forward(flip, TA, (A, TA), ([1, 3], Const); atol, rtol, fkwargs = (inv = true,)) + EnzymeTestUtils.test_forward(flip, TA, (A, TA), (1, Const); atol, rtol) + EnzymeTestUtils.test_forward(flip, TA, (A, TA), ([1, 3], Const); atol, rtol) end end end