Fix reverse rules: delegate via Enzyme.autodiff, return zeros for Active args

The reverse rules from #43 had two bugs exposed by new end-to-end tests:

1. Const-dret reverse returned `nothing` per arg, but Enzyme's rule
   protocol requires concrete scalar gradients for Active args (not
   nothing). Fixed to return `zero(T)` for Active args and `nothing`
   for Duplicated/Const args.

2. IIP reverse with Duplicated args (SciML pattern) returned nothing
   and never propagated gradients into the Duplicated shadow buffers.
   Fixed by delegating to `Enzyme.autodiff(Reverse, Const(f_orig),
   Const, args...)` when Duplicated args are present, so Enzyme
   accumulates the transposed derivative into the shadow buffers.

3. Enzyme passes `Type{<:Const}` (not an instance) for the dret slot
   in Const-return reverse rules. Updated dispatch signatures from
   `dret::EnzymeCore.Const` to `dret::Type{<:EnzymeCore.Const}`.

New end-to-end reverse-mode tests that assert derivative correctness:
- Const return + Active args: gradients are (0.0, 0.0)
- IIP f!(du, u) with Duplicated args: u_shadow accumulates ∂du/∂u
- Multi-component IIP cross-coupled Jacobian transpose
- ReverseWithPrimal IIP variant

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

ext/FunctionWrappersWrappersEnzymeExt.jl

@@ -210,26 +210,37 @@ function EnzymeRules.reverse(
     end
 end
 
-# Const return (no derivative to propagate from the return) — uniform Active args.
+# Const return — Enzyme passes the RT as a `Type{<:Const}` to `reverse`, not
+# as an instance.  Delegate the reverse pass to
+# `Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)` so gradients
+# accumulate into any `Duplicated` arg shadow buffers (the SciML IIP
+# pattern).  Simply returning `nothing` left Duplicated shadows at zero.
+#
+# Per Enzyme's rule return-type protocol, `Active` args require a concrete
+# scalar gradient (not `nothing`).  Under a `Const` return there is no
+# gradient source, so Active arg gradients are zero.  `Duplicated` /
+# `BatchDuplicated` args return `nothing` because their gradients are
+# accumulated in-place by the `Enzyme.autodiff(Reverse, …)` call above.
 function EnzymeRules.reverse(
     config::EnzymeRules.RevConfig,
     func::EnzymeCore.Const{<:FunctionWrappersWrapper},
-    dret::EnzymeCore.Const,
-    tape,
-    args::Vararg{EnzymeCore.Active, N}
-) where {N}
-    return ntuple(_ -> nothing, Val(N))
-end
-
-# Const return — mixed Active/Const args.
-function EnzymeRules.reverse(
-    config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
-    dret::EnzymeCore.Const,
+    dret::Type{<:EnzymeCore.Const},
     tape,
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {N}
-    return ntuple(_ -> nothing, Val(N))
+    f_orig = unwrap(func.val)
+    # Only worth invoking Enzyme.autodiff when at least one arg is
+    # Duplicated/BatchDuplicated — otherwise there's nothing to accumulate.
+    if any(a -> a isa EnzymeCore.Duplicated || a isa EnzymeCore.BatchDuplicated, args)
+        Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)
+    end
+    return ntuple(Val(N)) do i
+        if args[i] isa EnzymeCore.Active
+            zero(eltype(typeof(args[i])))
+        else
+            nothing
+        end
+    end
 end
 
 end

test/enzyme_tests.jl

@@ -132,12 +132,14 @@ end
     @test aug.shadow === nothing
     @test aug.tape === nothing
 
-    # Reverse step — dret is Const, no grads to accumulate.
+    # Reverse step — dret is Const (passed as TYPE not instance in reverse
+    # rules).  Enzyme's rule protocol requires concrete gradients for Active
+    # args; under a Const return they're zero (no gradient source).
     grads = EnzymeRules.reverse(
-        rconfig, Const(fww), EnzymeCore.Const{Float64}(0.0),
+        rconfig, Const(fww), EnzymeCore.Const{Float64},
         aug.tape, Active(3.0), Active(4.0)
     )
-    @test grads == (nothing, nothing)
+    @test grads == (0.0, 0.0)
 end
 
 @testset "Enzyme reverse mode, Duplicated return — augmented_primal initializes shadow" begin
@@ -171,3 +173,107 @@ end
     @test aug.tape === nothing
 end
 
+# =============================================================================
+# End-to-end reverse-mode derivative tests — exercise Enzyme.autodiff(Reverse,
+# …) through the FWW and assert the resulting gradients are numerically correct.
+# The prior reverse-mode testsets only checked dispatch / shape of
+# AugmentedReturn; they did NOT verify the gradients are right.
+# =============================================================================
+
+@testset "Enzyme Reverse: Const return, Active args — no-flow gradients" begin
+    # For a function whose return is annotated Const in Reverse mode, there is
+    # no gradient source from the return, so Active arg gradients must be 0.
+    # (Enzyme's rule-return protocol requires concrete gradients for Active
+    # args — `nothing` is not allowed — so the rule returns zeros.)
+    g(x, y) = x * y + x^2
+    fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,))
+
+    # Const return (instead of Active) → no gradient flows back
+    result = Enzyme.autodiff(Reverse, Const(fww), Const, Active(3.0), Active(4.0))
+    @test result[1] === (0.0, 0.0)
+end
+
+@testset "Enzyme Reverse: IIP with Duplicated args, Const return" begin
+    # SciML's standard pattern: IIP RHS `f!(du, u)` with Const return, both du
+    # and u are Duplicated.  Reverse mode should accumulate
+    #    u_shadow[i] += du_shadow[j] * ∂(du[j])/∂(u[i])
+    # into u_shadow.  For f!(du, u) = (du[1] = u[1]^2; nothing) with
+    #   du_shadow = [1.0] (incoming adjoint),
+    #   u[1] = 3.0,
+    #   ∂du[1]/∂u[1] = 2*u[1] = 6,
+    # the expected result is u_shadow[1] = 6.0 after the call.
+    f!(du, u) = (du[1] = u[1]^2; nothing)
+    fww = FunctionWrappersWrapper(
+        f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,)
+    )
+
+    du = [0.0]
+    u = [3.0]
+    du_shadow = [1.0]
+    u_shadow = [0.0]
+
+    Enzyme.autodiff(
+        Reverse, Const(fww), Const,
+        Duplicated(du, du_shadow), Duplicated(u, u_shadow)
+    )
+    @test du[1] ≈ 9.0          # primal effect: du[1] = u[1]^2 = 9
+    @test u_shadow[1] ≈ 6.0    # reverse accumulation: 2 * u[1] * du_shadow[1]
+end
+
+@testset "Enzyme Reverse: IIP multi-component IIP with Duplicated args" begin
+    # Cross-coupled IIP RHS: each output depends on multiple inputs.
+    #   du[1] = u[1] * u[2]
+    #   du[2] = u[1]^2 + u[2]^3
+    # Jacobian at u = (x, y):
+    #   J = [  y     x  ;
+    #         2x   3y^2 ]
+    # In reverse mode with du_shadow = [a, b], transpose of J applied to
+    # du_shadow gives the accumulation into u_shadow:
+    #   u_shadow[1] += a*y + b*2x
+    #   u_shadow[2] += a*x + b*3y^2
+    f!(du, u) = (du[1] = u[1]*u[2]; du[2] = u[1]^2 + u[2]^3; nothing)
+    fww = FunctionWrappersWrapper(
+        f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,)
+    )
+
+    x, y = 2.0, 5.0
+    a, b = 1.0, 0.5
+    du = zeros(2)
+    u = [x, y]
+    du_shadow = [a, b]
+    u_shadow = zeros(2)
+
+    Enzyme.autodiff(
+        Reverse, Const(fww), Const,
+        Duplicated(du, du_shadow), Duplicated(u, u_shadow)
+    )
+    @test du ≈ [x*y, x^2 + y^3]
+    @test u_shadow[1] ≈ a*y + b*2*x        # 5 + 2 = 7
+    @test u_shadow[2] ≈ a*x + b*3*y^2      # 2 + 37.5 = 39.5
+end
+
+@testset "Enzyme ReverseWithPrimal: IIP with Duplicated args" begin
+    # Same IIP pattern but with ReverseWithPrimal so we also check the primal
+    # is available when the rule is asked to include it.
+    f!(du, u) = (du[1] = u[1]^3; nothing)
+    fww = FunctionWrappersWrapper(
+        f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,)
+    )
+
+    du = [0.0]
+    u = [2.0]
+    du_shadow = [1.0]
+    u_shadow = [0.0]
+
+    # Capture the expected gradient BEFORE the call — Enzyme may zero
+    # `du_shadow` after consuming it during the reverse pass.
+    expected_u_grad = 3 * u[1]^2 * du_shadow[1]  # = 12.0
+
+    Enzyme.autodiff(
+        ReverseWithPrimal, Const(fww), Const,
+        Duplicated(du, du_shadow), Duplicated(u, u_shadow)
+    )
+    @test du[1] ≈ 8.0
+    @test u_shadow[1] ≈ expected_u_grad
+end
+