Add Mooncake extension for FunctionWrappersWrapper

When Mooncake differentiates through code that calls a
FunctionWrappersWrapper, it tries to create tangent types for each
FunctionWrapper variant in the internal tuple. These variants have
different type parameters (for different ForwardDiff Dual
combinations), producing incompatible FunctionWrapperTangent types
that can't be stored in a typed tuple — causing a convert error.

The fix: make FunctionWrappersWrapper calls a Mooncake primitive that
unwraps to the original function (via `unwrap`) and differentiates
through that directly. This mirrors the existing Enzyme extension
pattern. The FunctionWrappersWrapper itself gets NoTangent since it's
runtime dispatch infrastructure, not differentiable data — the
original function's derivatives are handled in the rrule.

This enables Mooncake to differentiate through NonlinearProblem solves
that use AutoSpecialize (FunctionWrappers), which is needed for
SCCNonlinearProblem AD support in SciMLSensitivity.jl (#1358).

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

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "FunctionWrappersWrappers"
 uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 authors = ["Chris Elrod <elrodc@gmail.com> and contributors"]
-version = "1.1.0"
+version = "1.2.0"
 
 [deps]
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
@@ -11,23 +11,27 @@ TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 [weakdeps]
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
 FunctionWrappersWrappersEnzymeExt = ["Enzyme", "EnzymeCore"]
+FunctionWrappersWrappersMooncakeExt = "Mooncake"
 
 [compat]
 Enzyme = "0.13"
 EnzymeCore = "0.8"
 FunctionWrappers = "1"
+Mooncake = "0.5"
 PrecompileTools = "1"
 TruncatedStacktraces = "1"
 julia = "1.10"
 
 [extras]
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Pkg", "Test", "Enzyme", "EnzymeCore"]
+test = ["Pkg", "Test", "Enzyme", "EnzymeCore", "Mooncake"]

ext/FunctionWrappersWrappersMooncakeExt.jl

@@ -0,0 +1,29 @@
+module FunctionWrappersWrappersMooncakeExt
+
+using FunctionWrappersWrappers
+import Mooncake
+using Mooncake: @is_primitive, MinimalCtx, CoDual, NoRData, zero_tangent, NoTangent
+
+# 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.
+
+@is_primitive MinimalCtx Tuple{<:FunctionWrappersWrapper, Vararg}
+
+function Mooncake.rrule!!(
+        f::CoDual{<:FunctionWrappersWrapper}, args::Vararg{CoDual},
+    )
+    f_orig = unwrap(f.x)
+    f_orig_codual = CoDual(f_orig, zero_tangent(f_orig))
+    y, pb = Mooncake.rrule!!(f_orig_codual, args...)
+    fww_pb(dy) = (NoRData(), Mooncake.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.
+Mooncake.tangent_type(::Type{<:FunctionWrappersWrapper}) = NoTangent
+
+end