Merge pull request #39 from ChrisRackauckas-Claude/mooncake-extension

Add Mooncake extension for FunctionWrappersWrapper

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

test/mooncake_tests.jl

@@ -0,0 +1,77 @@
+using FunctionWrappersWrappers
+using Mooncake
+using Test
+
+@testset "Mooncake reverse mode - single arg" begin
+    f(x) = x^2
+    fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,))
+
+    rule = Mooncake.build_rrule(fww, 3.0)
+    val, (_, dx) = Mooncake.value_and_gradient!!(rule, fww, 3.0)
+    @test val ≈ 9.0
+    @test dx ≈ 6.0
+end
+
+@testset "Mooncake reverse mode - multi arg" begin
+    g(x, y) = x * y + x^2
+    fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,))
+
+    # g(x,y) = x*y + x^2 → ∂g/∂x = y + 2x, ∂g/∂y = x
+    rule = Mooncake.build_rrule(fww, 3.0, 4.0)
+    val, (_, dx, dy) = Mooncake.value_and_gradient!!(rule, fww, 3.0, 4.0)
+    @test val ≈ 21.0
+    @test dx ≈ 10.0  # ∂g/∂x at (3,4) = 4 + 6
+    @test dy ≈ 3.0   # ∂g/∂y at (3,4) = 3
+end
+
+@testset "Mooncake with trig functions" begin
+    fww_sin = FunctionWrappersWrapper(sin, (Tuple{Float64},), (Float64,))
+
+    rule = Mooncake.build_rrule(fww_sin, 1.0)
+    val, (_, dx) = Mooncake.value_and_gradient!!(rule, fww_sin, 1.0)
+    @test val ≈ sin(1.0)
+    @test dx ≈ cos(1.0)
+end
+
+@testset "Mooncake through loss function" begin
+    # Test that Mooncake can differentiate a loss function that calls FunctionWrappersWrapper
+    f(x) = x[1]^2 + x[2]^2
+    fww = FunctionWrappersWrapper(f, (Tuple{Vector{Float64}},), (Float64,))
+
+    loss(x) = fww(x)
+    rule = Mooncake.build_rrule(loss, [3.0, 4.0])
+    val, (_, dx) = Mooncake.value_and_gradient!!(rule, loss, [3.0, 4.0])
+    @test val ≈ 25.0
+    @test dx ≈ [6.0, 8.0]
+end
+
+@testset "Mooncake in-place function" begin
+    # In-place functions are common in SciML (f!(du, u, p, t))
+    function f!(du, u, p)
+        du[1] = p[1] * u[1] + p[2] * u[2]
+        du[2] = p[3] * u[1] - u[2]
+        return nothing
+    end
+    fww = FunctionWrappersWrapper(
+        f!,
+        (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}},),
+        (Nothing,),
+    )
+
+    function loss(p)
+        u = [1.0, 2.0]
+        du = similar(u)
+        fww(du, u, p)
+        return sum(abs2, du)
+    end
+
+    rule = Mooncake.build_rrule(loss, [1.0, 2.0, 3.0])
+    val, (_, dp) = Mooncake.value_and_gradient!!(rule, loss, [1.0, 2.0, 3.0])
+    # f!(du, [1,2], [1,2,3]) → du = [1*1+2*2, 3*1-2] = [5, 1]
+    # loss = 25 + 1 = 26
+    @test val ≈ 26.0
+    # ∂loss/∂p1 = 2*du[1]*u[1] = 2*5*1 = 10
+    # ∂loss/∂p2 = 2*du[1]*u[2] = 2*5*2 = 20
+    # ∂loss/∂p3 = 2*du[2]*u[1] = 2*1*1 = 2
+    @test dp ≈ [10.0, 20.0, 2.0]
+end

test/runtests.jl

@@ -59,3 +59,9 @@ if GROUP == "All" || GROUP == "Enzyme"
         include("enzyme_tests.jl")
     end
 end
+
+if GROUP == "All" || GROUP == "Mooncake"
+    @testset "Mooncake extension" begin
+        include("mooncake_tests.jl")
+    end
+end