Add forward mode rules for calling a FunctionWrappersWrapper

ext/FunctionWrappersWrappersMooncakeExt.jl

@@ -2,35 +2,96 @@ module FunctionWrappersWrappersMooncakeExt
 
 using FunctionWrappersWrappers
 import Mooncake
-using Mooncake: @is_primitive, MinimalCtx, CoDual, NoRData, zero_tangent, NoTangent, fdata
+using Mooncake:
+    @is_primitive,
+    MinimalCtx,
+    CoDual,
+    Dual,
+    NoRData,
+    zero_tangent,
+    NoTangent,
+    fdata,
+    primal,
+    build_frule
 
-# Make calling a FunctionWrappersWrapper a Mooncake primitive.
-# Instead of differentiating through the FunctionWrapper dispatch machinery
-# (which fails because the tuple of differently-typed FunctionWrappers produces
-# incompatible FunctionWrapperTangent types), unwrap to the original function
-# and differentiate through that directly.
+# Unwrap to the original function and differentiate through that, rather than the
+# FunctionWrapper dispatch machinery, which fails on mismatched FunctionWrapperTangent
+# types when the tuple holds differently-typed wrappers.
 
 @is_primitive MinimalCtx Tuple{<:FunctionWrappersWrapper, Vararg}
 
+# unwrap reaches into FunctionWrapper's internal .obj field directly, which the generic
+# getfield rule can't handle since FunctionWrapperTangent isn't a StandardTangentType.
+# Also needed for HVP: forward-mode has to trace through this call when differentiating
+# the rrule!!/frule!! above.
+@is_primitive MinimalCtx Tuple{typeof(unwrap), <:FunctionWrappersWrapper}
+
+# Every rule below builds a fresh zero tangent/dual for the unwrapped function rather than
+# threading through any tangent it might really carry. That's exact only if the unwrapped
+# function is genuinely non-differentiable (e.g. a stateless closure); if it carries real
+# state (e.g. a struct field holding a learnable Matrix), silently zeroing it would give a
+# wrong, too-small gradient instead of an error. Fail loud instead.
+function _check_wrapped_fn_has_no_tangent(f_orig)
+    Mooncake.tangent_type(typeof(f_orig)) === NoTangent && return nothing
+    return error(
+        "Differentiating through a FunctionWrappersWrapper whose wrapped function " *
+            "itself carries differentiable state (tangent_type = " *
+            "$(Mooncake.tangent_type(typeof(f_orig)))) is not supported: only " *
+            "gradients flowing through the call arguments are tracked here, not " *
+            "through the wrapped function's own fields.",
+    )
+end
+
+function Mooncake.rrule!!(::CoDual{typeof(unwrap)}, fww::CoDual{<:FunctionWrappersWrapper})
+    f_orig = unwrap(fww.x)
+    _check_wrapped_fn_has_no_tangent(f_orig)
+    unwrap_pb(::NoRData) = (NoRData(), NoRData())
+    return CoDual(f_orig, fdata(zero_tangent(f_orig))), unwrap_pb
+end
+
+function Mooncake.frule!!(::Dual{typeof(unwrap)}, fww::Dual{<:FunctionWrappersWrapper})
+    f_orig = unwrap(primal(fww))
+    _check_wrapped_fn_has_no_tangent(f_orig)
+    return Dual(f_orig, zero_tangent(f_orig))
+end
+
+# Cache derived rules by signature: f_orig gets called once per ODE timestep, so rebuilding
+# one from scratch on every call would be wasted work once the signature has stabilised.
+const _CALL_RRULE_CACHE = Dict{Any, Any}()
+const _CALL_FRULE_CACHE = Dict{Any, Any}()
+
 function Mooncake.rrule!!(
         f::CoDual{<:FunctionWrappersWrapper}, args::Vararg{CoDual},
     )
     f_orig = unwrap(f.x)
-    # Build a derived rule for calling the unwrapped function with these arg types.
-    # We can't use rrule!! directly since the unwrapped function (e.g. SciMLBase.Void)
-    # is generally not a Mooncake primitive — it needs a derived (compiled) rule.
+    _check_wrapped_fn_has_no_tangent(f_orig)
+    # The unwrapped function usually isn't a Mooncake primitive, so build a derived rule.
     sig = Tuple{typeof(f_orig), map(Core.Typeof ∘ Mooncake.primal, args)...}
-    rule = Mooncake.build_rrule(sig)
-    # Use fdata to get the correct tangent component for the CoDual — zero_tangent
-    # returns NoTangent for singleton callables but derived rules expect NoFData.
+    rule = get!(() -> Mooncake.build_rrule(sig), _CALL_RRULE_CACHE, sig)
+    # fdata turns zero_tangent's NoTangent into the NoFData a derived rule expects.
     f_orig_codual = CoDual(f_orig, fdata(zero_tangent(f_orig)))
     y, pb = rule(f_orig_codual, args...)
     fww_pb(dy) = (NoRData(), Base.tail(pb(dy))...)
     return y, fww_pb
 end
 
-# FunctionWrappersWrapper is not differentiable data itself — the wrapped function
-# is what carries the derivative information, and we handle that in the rrule above.
+function Mooncake.frule!!(
+        f::Dual{<:FunctionWrappersWrapper}, args::Vararg{Dual},
+    )
+    f_orig = unwrap(primal(f))
+    _check_wrapped_fn_has_no_tangent(f_orig)
+    f_orig_dual = Dual(f_orig, zero_tangent(f_orig))
+    # Mirrors the rrule!! above, but builds a derived forward-mode rule instead. `sig` is
+    # only used as the cache key here; the rule itself is built via `build_frule`'s
+    # args-based convenience method (mirroring `build_rrule(sig)`'s use above, just with
+    # actual values instead of a signature type -- this form doesn't need `get_interpreter`
+    # /`ForwardMode` imported at all, since it extracts the interpreter internally).
+    sig = Tuple{typeof(f_orig), map(Core.Typeof ∘ primal, args)...}
+    rule = get!(() -> build_frule(f_orig_dual, args...), _CALL_FRULE_CACHE, sig)
+    return rule(f_orig_dual, args...)
+end
+
+# The wrapper itself carries no derivative info; the wrapped function does, handled above.
 Mooncake.tangent_type(::Type{<:FunctionWrappersWrapper}) = NoTangent
 
 end

test/qa/qa.jl

@@ -21,10 +21,10 @@ run_qa(
         # (`value_and_gradient!!`, `prepare_gradient_cache`, `Config`, `Dual`, ...).
         # Its whole rule-authoring interface — the thing an extension that teaches
         # Mooncake about a new callable has to use — is documented but not `public`:
-        # `build_rrule`, `rrule!!`, `tangent_type` and `primal` below, plus the
-        # explicitly-imported names in the matching list. There is no public spelling
-        # for writing a Mooncake rule, so these stay ignored until Mooncake marks its
-        # rule API `public`.
+        # `build_rrule`, `rrule!!`, `frule!!`, `tangent_type` and `primal` below, plus
+        # the explicitly-imported names in the matching list. There is no public
+        # spelling for writing a Mooncake rule, so these stay ignored until Mooncake
+        # marks its rule API `public`.
         #
         # `Core.Typeof` is the standard idiom for building a call-signature tuple type
         # (`typeof` is wrong for arguments that are themselves types); `Core` does not
@@ -40,14 +40,14 @@ run_qa(
             ignore = (
                 :FunctionWrapper,
                 :augmented_primal, :forward, :inactive_type, :reverse, :strong_zero,
-                :build_rrule, :primal, Symbol("rrule!!"), :tangent_type,
+                :build_rrule, :primal, Symbol("rrule!!"), Symbol("frule!!"), :tangent_type,
                 :Typeof,
             ),
         ),
         all_explicit_imports_are_public = (;
             ignore = (
                 Symbol("@is_primitive"), :CoDual, :MinimalCtx, :NoRData, :NoTangent,
-                :fdata, :zero_tangent,
+                :fdata, :zero_tangent, :build_frule, :primal,
             ),
         ),
     ),