Merge pull request #49 from ChrisRackauckas-Claude/fix-issue-48-duplicated-fww

Accept Annotation{<:FunctionWrappersWrapper} in Enzyme rules (closes #48)

Author: ChrisRackauckas

ext/FunctionWrappersWrappersEnzymeExt.jl

@@ -36,9 +36,19 @@ end
 # =============================================================================
 
 # Shadow only (Forward mode, no primal)
+#
+# `func` is `Annotation{<:FunctionWrappersWrapper}` rather than
+# `Const{<:FunctionWrappersWrapper}` so that callers passing
+# `Duplicated{<:FunctionWrappersWrapper}` also dispatch here.  Enzyme drives
+# the rule that way when the outer `autodiff` call is differentiating through
+# a closure that carries an FWW (e.g. NonlinearSolve + SciMLSensitivity, see
+# SciML/FunctionWrappersWrappers.jl#48).  The FWW struct itself only carries
+# `FunctionWrapper`s plus cache storage — none of those fields have a
+# meaningful tangent — so the function shadow is ignored and the inner
+# `Enzyme.autodiff` call uses `Const(f_orig)`.
 function EnzymeRules.forward(
     ::EnzymeRules.FwdConfig{false, true, W, RuntimeActivity, StrongZero},
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation{T}},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {T, W, N, RuntimeActivity, StrongZero}
@@ -56,7 +66,7 @@ end
 # Both primal and shadow (ForwardWithPrimal mode)
 function EnzymeRules.forward(
     ::EnzymeRules.FwdConfig{true, true, W, RuntimeActivity, StrongZero},
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation{T}},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {T, W, N, RuntimeActivity, StrongZero}
@@ -81,7 +91,7 @@ end
 # Primal only (Const return type) — width-independent
 function EnzymeRules.forward(
     ::EnzymeRules.FwdConfig{true, false, W, RuntimeActivity, StrongZero},
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {W, N, RuntimeActivity, StrongZero}
@@ -107,7 +117,7 @@ end
 # `set_runtime_activity(Forward)` on the way down into `f_orig`. 
 function EnzymeRules.forward(
     ::EnzymeRules.FwdConfig{false, false, W, RuntimeActivity, StrongZero},
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Annotation},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {W, N, RuntimeActivity, StrongZero}
@@ -123,7 +133,7 @@ end
 
 function EnzymeRules.augmented_primal(
     config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Active{T}},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {T, N}
@@ -143,7 +153,7 @@ end
 # the reverse pass has nothing to propagate back from the return.
 function EnzymeRules.augmented_primal(
     config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Const},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {N}
@@ -157,7 +167,7 @@ end
 # it available when propagating dret through the arguments.
 function EnzymeRules.augmented_primal(
     config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.Duplicated{T}},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {T, N}
@@ -173,7 +183,7 @@ end
 
 function EnzymeRules.augmented_primal(
     config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     RT::Type{<:EnzymeCore.BatchDuplicated{T, W}},
     args::Vararg{EnzymeCore.Annotation, N}
 ) where {T, W, N}
@@ -220,7 +230,7 @@ end
 
 function EnzymeRules.reverse(
     config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     dret::EnzymeCore.Active{T},
     tape,
     args::Vararg{EnzymeCore.Active, N}
@@ -232,7 +242,7 @@ end
 # Handle mixed Active/Const args: return nothing for Const, gradient for Active
 function EnzymeRules.reverse(
     config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     dret::EnzymeCore.Active,
     tape,
     args::Vararg{EnzymeCore.Annotation, N}
@@ -271,7 +281,7 @@ end
 # accumulated in-place by the `Enzyme.autodiff(Reverse, …)` call above.
 function EnzymeRules.reverse(
     config::EnzymeRules.RevConfig,
-    func::EnzymeCore.Const{<:FunctionWrappersWrapper},
+    func::EnzymeCore.Annotation{<:FunctionWrappersWrapper},
     dret::Type{<:EnzymeCore.Const},
     tape,
     args::Vararg{EnzymeCore.Annotation, N}

test/enzyme_tests.jl

@@ -298,6 +298,112 @@ end
 # numerically correct.
 # =============================================================================
 
+# =============================================================================
+# Duplicated function annotation on the FWW itself.
+#
+# Reproduces SciML/FunctionWrappersWrappers.jl#48: when Enzyme differentiates
+# through a closure that captures an FWW (e.g. NonlinearSolve +
+# SciMLSensitivity), the rule is invoked with
+# `Duplicated{<:FunctionWrappersWrapper}` for the function argument, not
+# `Const{<:FunctionWrappersWrapper}`.  The FWW struct itself only carries
+# `FunctionWrapper`s + cache storage, so its "shadow" is ignored — we route
+# through `unwrap(func.val)` exactly as with `Const`.
+# =============================================================================
+
+@testset "Enzyme forward, Duplicated FWW annotation — IIP Const return" begin
+    f!(residual, u, p) = (residual[1] = u[1]^2 - p[1]; nothing)
+    fww = FunctionWrappersWrapper(
+        f!,
+        (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}},),
+        (Nothing,)
+    )
+
+    residual = [0.0]; dresidual = [0.0]
+    u = [2.0];        du = [1.0]
+    p = [1.0];        dp = [0.0]
+
+    config = EnzymeCore.EnzymeRules.FwdConfig{false, false, 1, false, false}()
+    ret = EnzymeCore.EnzymeRules.forward(
+        config,
+        Duplicated(fww, fww),                # <-- the failing dispatch in #48
+        EnzymeCore.Const{Nothing},
+        Duplicated(residual, dresidual),
+        Duplicated(u, du),
+        Duplicated(p, dp),
+    )
+    @test ret === nothing
+    @test residual[1] ≈ 3.0      # u[1]^2 - p[1] = 4 - 1
+    @test dresidual[1] ≈ 4.0     # 2*u[1]*du[1] - 1*dp[1] = 4
+end
+
+@testset "Enzyme forward, Duplicated FWW annotation — shadow-only return" begin
+    # Drive the {false, true, W, …} rule (shadow only, no primal) with a
+    # Duplicated FWW.
+    f(x) = x^2
+    fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,))
+
+    config = EnzymeCore.EnzymeRules.FwdConfig{false, true, 1, false, false}()
+    shadow = EnzymeCore.EnzymeRules.forward(
+        config,
+        Duplicated(fww, fww),
+        EnzymeCore.Duplicated{Float64},
+        Duplicated(3.0, 1.0),
+    )
+    @test shadow ≈ 6.0           # f'(3) = 2*3 = 6
+end
+
+@testset "Enzyme forward, Duplicated FWW annotation — primal + shadow return" begin
+    # Drive the {true, true, W, …} rule (ForwardWithPrimal) with a Duplicated
+    # FWW.
+    f(x) = x^2
+    fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,))
+
+    config = EnzymeCore.EnzymeRules.FwdConfig{true, true, 1, false, false}()
+    result = EnzymeCore.EnzymeRules.forward(
+        config,
+        Duplicated(fww, fww),
+        EnzymeCore.Duplicated{Float64},
+        Duplicated(3.0, 1.0),
+    )
+    @test result isa Duplicated
+    @test result.val ≈ 9.0       # primal
+    @test result.dval ≈ 6.0      # shadow
+end
+
+@testset "Enzyme reverse, Duplicated FWW annotation — Const return IIP" begin
+    # Mirror the forward IIP case on the reverse side.  Duplicated FWW must
+    # still drive the rule, gradients must accumulate into u_shadow.
+    f!(du, u) = (du[1] = u[1]^2; nothing)
+    fww = FunctionWrappersWrapper(
+        f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,)
+    )
+
+    du = [0.0];       du_shadow = [1.0]
+    u  = [3.0];       u_shadow  = [0.0]
+
+    rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false), false, false}()
+    aug = EnzymeRules.augmented_primal(
+        rconfig,
+        Duplicated(fww, fww),                # <-- Duplicated FWW
+        EnzymeCore.Const{Nothing},
+        Duplicated(du, du_shadow),
+        Duplicated(u, u_shadow),
+    )
+    @test aug.primal === nothing
+    @test aug.shadow === nothing
+
+    EnzymeRules.reverse(
+        rconfig,
+        Duplicated(fww, fww),
+        EnzymeCore.Const{Nothing},
+        aug.tape,
+        Duplicated(du, du_shadow),
+        Duplicated(u, u_shadow),
+    )
+    @test du[1] ≈ 9.0            # primal effect from augmented_primal
+    @test u_shadow[1] ≈ 6.0      # reverse accumulation: 2*u[1]*du_shadow[1]
+end
+
 @testset "Enzyme Forward: set_runtime_activity propagates through FWW (IIP, time-dependent)" begin
     # DiffEqBase's `wrapfun_iip(ff, (u, u, p, t))` shape.
     const_INPUTS = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64}