diff --git a/ext/FunctionWrappersWrappersMooncakeExt.jl b/ext/FunctionWrappersWrappersMooncakeExt.jl index 3e2aa17..2d423c8 100644 --- a/ext/FunctionWrappersWrappersMooncakeExt.jl +++ b/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 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, ), ), ),