From 94e4164cad07e855124edb9e84cb4dd5751c4cd8 Mon Sep 17 00:00:00 2001 From: AstitvaAggarwal Date: Mon, 10 Aug 2026 12:38:16 +0100 Subject: [PATCH 1/4] Add forward-mode rules for FunctionWrappersWrapper calls and unwrap --- ext/FunctionWrappersWrappersMooncakeExt.jl | 61 +++++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/ext/FunctionWrappersWrappersMooncakeExt.jl b/ext/FunctionWrappersWrappersMooncakeExt.jl index 3e2aa17..7cc2d7a 100644 --- a/ext/FunctionWrappersWrappersMooncakeExt.jl +++ b/ext/FunctionWrappersWrappersMooncakeExt.jl @@ -2,7 +2,19 @@ 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, + get_interpreter, + ForwardMode # Make calling a FunctionWrappersWrapper a Mooncake primitive. # Instead of differentiating through the FunctionWrapper dispatch machinery @@ -12,6 +24,30 @@ using Mooncake: @is_primitive, MinimalCtx, CoDual, NoRData, zero_tangent, NoTang @is_primitive MinimalCtx Tuple{<:FunctionWrappersWrapper, Vararg} +# `unwrap` reaches directly into a nested `FunctionWrapper`'s internal `.obj` field +# (bypassing the call/construct interface entirely), which the generic `lgetfield` rule +# cannot handle here: it's constrained to `StandardTangentType`/`StandardFDataType`, and +# `FunctionWrapper`'s custom `FunctionWrapperTangent` isn't one. This matters for +# forward-over-reverse (HVP): computing a derivative of `rrule!!`/`frule!!` above requires +# tracing their own bodies, including this `unwrap` call, in forward mode. Since we already +# always rebuild a fresh `zero_tangent` for the unwrapped function in the rules above +# (never actually reading any tangent carried by the wrapper itself, consistent with +# `FunctionWrappersWrapper`'s own tangent being `NoTangent`), treating `unwrap` as an +# opaque primitive with the same "fresh zero tangent" behaviour is exact, not an +# approximation. +@is_primitive MinimalCtx Tuple{typeof(unwrap), <:FunctionWrappersWrapper} + +function Mooncake.rrule!!(::CoDual{typeof(unwrap)}, fww::CoDual{<:FunctionWrappersWrapper}) + f_orig = unwrap(fww.x) + 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)) + return Dual(f_orig, zero_tangent(f_orig)) +end + function Mooncake.rrule!!( f::CoDual{<:FunctionWrappersWrapper}, args::Vararg{CoDual}, ) @@ -29,8 +65,31 @@ function Mooncake.rrule!!( return y, fww_pb end +function Mooncake.frule!!( + f::Dual{<:FunctionWrappersWrapper}, args::Vararg{Dual}, + ) + f_orig = unwrap(primal(f)) + # Mirrors the rrule!! above: build a derived forward-mode rule for the unwrapped + # function with these arg types, rather than differentiating through the + # FunctionWrappersWrapper dispatch machinery itself. + sig = Tuple{typeof(f_orig), map(Core.Typeof ∘ primal, args)...} + rule = build_frule(get_interpreter(ForwardMode), sig) + f_orig_dual = Dual(f_orig, zero_tangent(f_orig)) + return rule(f_orig_dual, args...) +end + # FunctionWrappersWrapper is not differentiable data itself — the wrapped function # is what carries the derivative information, and we handle that in the rrule above. Mooncake.tangent_type(::Type{<:FunctionWrappersWrapper}) = NoTangent +# For the same reason, `prepare_pullback_cache`/`value_and_pullback!!`'s generic +# "no pointers or aliased mutable state reachable from the output" safety check +# (SciML/SciMLSensitivity.jl#1424) has nothing to protect against here either: its +# `.fw` field holds raw-`Ptr`-carrying `FunctionWrapper`s and its `.cache_storage` +# field is mutable, shared cache state that can legitimately be aliased across +# multiple places in a returned value (e.g. an `ODESolution`) — but neither is ever +# reached via generic field access during real differentiation, only via the +# dedicated `rrule!!`/`unwrap` path above. Stop the check's recursion here. +Mooncake.__exclude_unsupported_output_internal!(::FunctionWrappersWrapper, ::Set{UInt}) = nothing + end From 34ea2c6dee41a12977baf952e305f7d9d9dd2197 Mon Sep 17 00:00:00 2001 From: AstitvaAggarwal Date: Mon, 10 Aug 2026 14:27:03 +0100 Subject: [PATCH 2/4] Guard against dropping the wrapped function's own gradient, cache derived rules --- ext/FunctionWrappersWrappersMooncakeExt.jl | 78 +++++++++++----------- 1 file changed, 40 insertions(+), 38 deletions(-) diff --git a/ext/FunctionWrappersWrappersMooncakeExt.jl b/ext/FunctionWrappersWrappersMooncakeExt.jl index 7cc2d7a..ff5ae8f 100644 --- a/ext/FunctionWrappersWrappersMooncakeExt.jl +++ b/ext/FunctionWrappersWrappersMooncakeExt.jl @@ -16,49 +16,63 @@ using Mooncake: get_interpreter, ForwardMode -# 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 directly into a nested `FunctionWrapper`'s internal `.obj` field -# (bypassing the call/construct interface entirely), which the generic `lgetfield` rule -# cannot handle here: it's constrained to `StandardTangentType`/`StandardFDataType`, and -# `FunctionWrapper`'s custom `FunctionWrapperTangent` isn't one. This matters for -# forward-over-reverse (HVP): computing a derivative of `rrule!!`/`frule!!` above requires -# tracing their own bodies, including this `unwrap` call, in forward mode. Since we already -# always rebuild a fresh `zero_tangent` for the unwrapped function in the rules above -# (never actually reading any tangent carried by the wrapper itself, consistent with -# `FunctionWrappersWrapper`'s own tangent being `NoTangent`), treating `unwrap` as an -# opaque primitive with the same "fresh zero tangent" behaviour is exact, not an -# approximation. +# 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, mirroring Mooncake's own DynamicRRule/DynamicFRule +# (src/interpreter/{reverse,forward}_mode.jl): f_orig gets called once per ODE timestep, so +# rebuilding (and re-locking Mooncake's internal rule cache) 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))...) @@ -69,27 +83,15 @@ function Mooncake.frule!!( f::Dual{<:FunctionWrappersWrapper}, args::Vararg{Dual}, ) f_orig = unwrap(primal(f)) - # Mirrors the rrule!! above: build a derived forward-mode rule for the unwrapped - # function with these arg types, rather than differentiating through the - # FunctionWrappersWrapper dispatch machinery itself. + _check_wrapped_fn_has_no_tangent(f_orig) + # Mirrors the rrule!! above, but builds a derived forward-mode rule instead. sig = Tuple{typeof(f_orig), map(Core.Typeof ∘ primal, args)...} - rule = build_frule(get_interpreter(ForwardMode), sig) + rule = get!(() -> build_frule(get_interpreter(ForwardMode), sig), _CALL_FRULE_CACHE, sig) f_orig_dual = Dual(f_orig, zero_tangent(f_orig)) return rule(f_orig_dual, args...) end -# FunctionWrappersWrapper is not differentiable data itself — the wrapped function -# is what carries the derivative information, and we handle that in the rrule above. +# The wrapper itself carries no derivative info; the wrapped function does, handled above. Mooncake.tangent_type(::Type{<:FunctionWrappersWrapper}) = NoTangent -# For the same reason, `prepare_pullback_cache`/`value_and_pullback!!`'s generic -# "no pointers or aliased mutable state reachable from the output" safety check -# (SciML/SciMLSensitivity.jl#1424) has nothing to protect against here either: its -# `.fw` field holds raw-`Ptr`-carrying `FunctionWrapper`s and its `.cache_storage` -# field is mutable, shared cache state that can legitimately be aliased across -# multiple places in a returned value (e.g. an `ODESolution`) — but neither is ever -# reached via generic field access during real differentiation, only via the -# dedicated `rrule!!`/`unwrap` path above. Stop the check's recursion here. -Mooncake.__exclude_unsupported_output_internal!(::FunctionWrappersWrapper, ::Set{UInt}) = nothing - end From 38d1c17dfc15a85a2e4e2ab5cff28320e9ebb069 Mon Sep 17 00:00:00 2001 From: AstitvaAggarwal Date: Mon, 10 Aug 2026 14:46:11 +0100 Subject: [PATCH 3/4] Trim internal-plumbing detail from cache comment --- ext/FunctionWrappersWrappersMooncakeExt.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/ext/FunctionWrappersWrappersMooncakeExt.jl b/ext/FunctionWrappersWrappersMooncakeExt.jl index ff5ae8f..304e479 100644 --- a/ext/FunctionWrappersWrappersMooncakeExt.jl +++ b/ext/FunctionWrappersWrappersMooncakeExt.jl @@ -57,10 +57,8 @@ function Mooncake.frule!!(::Dual{typeof(unwrap)}, fww::Dual{<:FunctionWrappersWr return Dual(f_orig, zero_tangent(f_orig)) end -# Cache derived rules by signature, mirroring Mooncake's own DynamicRRule/DynamicFRule -# (src/interpreter/{reverse,forward}_mode.jl): f_orig gets called once per ODE timestep, so -# rebuilding (and re-locking Mooncake's internal rule cache) on every call would be wasted -# work once the signature has stabilised. +# 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}() From 72a1b209392b1d2cdb5718e6fe2e4753d91aaad2 Mon Sep 17 00:00:00 2001 From: AstitvaAggarwal Date: Wed, 12 Aug 2026 18:05:05 +0100 Subject: [PATCH 4/4] Drop ForwardMode/get_interpreter, use build_frule's values-based form build_frule(f_orig_dual, args...) is an existing convenience method that extracts the interpreter/signature internally, same as build_rrule(sig) already does for the reverse-mode rule above. Removes the need to import ForwardMode/get_interpreter at all, cutting the Mooncake QA ignore list from 15 to 13 symbols. --- ext/FunctionWrappersWrappersMooncakeExt.jl | 14 ++++++++------ test/qa/qa.jl | 12 ++++++------ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/ext/FunctionWrappersWrappersMooncakeExt.jl b/ext/FunctionWrappersWrappersMooncakeExt.jl index 304e479..2d423c8 100644 --- a/ext/FunctionWrappersWrappersMooncakeExt.jl +++ b/ext/FunctionWrappersWrappersMooncakeExt.jl @@ -12,9 +12,7 @@ using Mooncake: NoTangent, fdata, primal, - build_frule, - get_interpreter, - ForwardMode + build_frule # Unwrap to the original function and differentiate through that, rather than the # FunctionWrapper dispatch machinery, which fails on mismatched FunctionWrapperTangent @@ -82,10 +80,14 @@ function Mooncake.frule!!( ) f_orig = unwrap(primal(f)) _check_wrapped_fn_has_no_tangent(f_orig) - # Mirrors the rrule!! above, but builds a derived forward-mode rule instead. - sig = Tuple{typeof(f_orig), map(Core.Typeof ∘ primal, args)...} - rule = get!(() -> build_frule(get_interpreter(ForwardMode), sig), _CALL_FRULE_CACHE, sig) 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 diff --git a/test/qa/qa.jl b/test/qa/qa.jl index 9f853a1..8ed505e 100644 --- a/test/qa/qa.jl +++ b/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, ), ), ),