diff --git a/ext/AbstractFastHartleyTransformsEnzymeCoreExt.jl b/ext/AbstractFastHartleyTransformsEnzymeCoreExt.jl index c3191a2..df79e5a 100644 --- a/ext/AbstractFastHartleyTransformsEnzymeCoreExt.jl +++ b/ext/AbstractFastHartleyTransformsEnzymeCoreExt.jl @@ -61,6 +61,26 @@ function EnzymeRules.forward( return nothing end +function EnzymeRules.forward( + config::FwdConfigWidth{1}, + ::Const{typeof(*)}, + RT::Type, + p::Const{<:DHTPlan}, + x::Const{<:Array}, + ) + primal = p.val * x.val + tangent = zero(primal) + + if needs_primal(config) && needs_shadow(config) + return Duplicated(primal, tangent) + elseif needs_primal(config) + return primal + elseif needs_shadow(config) + return tangent + end + return nothing +end + function EnzymeRules.augmented_primal( config::RevConfigWidth{1}, ::Const{typeof(*)}, @@ -163,4 +183,32 @@ function EnzymeRules.reverse( return (nothing, nothing, nothing) end +# Constant input: no cotangent needs to propagate to x. +function EnzymeRules.augmented_primal( + config::RevConfigWidth{1}, + ::Const{typeof(*)}, + RT::Type, + p::Const{<:DHTPlan}, + x::Const{<:Array}, + ) + primal = p.val * x.val + shadow = needs_shadow(config) ? zero(primal) : nothing + return AugmentedReturn( + needs_primal(config) ? primal : nothing, + shadow, + nothing, + ) +end + +function EnzymeRules.reverse( + config::RevConfigWidth{1}, + ::Const{typeof(*)}, + RT::Type, + tape, + p::Const{<:DHTPlan}, + x::Const{<:Array}, + ) + return (nothing, nothing) +end + end diff --git a/test/enzyme.jl b/test/enzyme.jl index fe4b61f..36b0907 100644 --- a/test/enzyme.jl +++ b/test/enzyme.jl @@ -193,3 +193,50 @@ LinearAlgebra.mul!(y::Array, ::AliasingIdentityPlan, x::Array) = copyto!(y, x) check_enzyme_loss(a -> aliased_buffer_loss(p, a, target), x, expected) end end + +function constant_input_loss(a, p, x) + y = p * x + return sum(abs2, a .* y) +end + +@testset "Enzyme constant Hartley input" begin + for T in (Float32, Float64) + x = reshape(T.(sin.(1:12)), 3, 4) + a = reshape(T.(cos.(1:12)), 3, 4) + p = plan_fht(x, (1, 2)) + tolerance = T == Float32 ? 2.0f-5 : 2.0e-11 + + expected = 2 .* a .* abs2.(p * x) + + # Reverse mode + da = zero(a) + Enzyme.autodiff( + Reverse, + constant_input_loss, + Active, + Duplicated(a, da), + Const(p), + Const(x), + ) + @test isapprox(da, expected; rtol = tolerance, atol = tolerance) + + # Forward mode + direction = reshape(T.(sin.(13:24)), 3, 4) + result = Enzyme.autodiff( + Forward, + constant_input_loss, + Duplicated, + Duplicated(copy(a), copy(direction)), + Const(p), + Const(x), + ) + + expected_directional = sum(expected .* direction) + @test isapprox( + result[1], + expected_directional; + rtol = tolerance, + atol = tolerance, + ) + end +end