Fix Enzyme differentiation with constant Hartley inputs

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

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