Mooncake for NonlinearSolve's Adjoints. (#719)

* fix broken broken tests for Mooncake.

* minor fix.

---------

Co-authored-by: Christopher Rackauckas <accounts@chrisrackauckas.com>

Author: AstitvaAggarwal

lib/NonlinearSolveBase/ext/NonlinearSolveBaseChainRulesCoreExt.jl

@@ -6,7 +6,7 @@ using SciMLBase
 using SciMLBase: AbstractSensitivityAlgorithm
 
 import ChainRulesCore
-import ChainRulesCore: NoTangent
+import ChainRulesCore: NoTangent, Tangent
 
 function ChainRulesCore.frule(::typeof(NonlinearSolveBase.solve_up), prob,
         sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
@@ -19,13 +19,23 @@ function ChainRulesCore.frule(::typeof(NonlinearSolveBase.solve_up), prob,
 end
 
 function ChainRulesCore.rrule(::typeof(NonlinearSolveBase.solve_up), prob::AbstractNonlinearProblem,
-        sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
-        u0, p, args...; originator = SciMLBase.ChainRulesOriginator(),
-        kwargs...)
-    NonlinearSolveBase._solve_adjoint(
+    sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
+    u0, p, args...; originator = SciMLBase.ChainRulesOriginator(),
+    kwargs...)
+    primal, inner_thunking_pb = NonlinearSolveBase._solve_adjoint(
         prob, sensealg, u0, p,
         originator, args...;
         kwargs...)
+
+    # when using mooncake ∂sol would be a NamedTuple Tangent with cotangents of all the solution struct's fields.
+    # However the pullback for this rule - "steadystatebackpass" as defined in SciMLSensitivity/src/concrete_solve.jl/
+    # handles AD only when ∂sol is a ChainRulesCore.AbstractThunk object or a sol.u vector and similar data structures (not namedtuples).
+    # When using Mooncake, we pass in sol.u to inner_thunking_pb directly as this is the only field relevant to the solution's cotangent (given solve_up, AbstractNonlinearProblem setting).
+
+    function solve_up_adjoint(∂sol)
+        return inner_thunking_pb(∂sol isa Tangent{Any,<:NamedTuple} ? ∂sol.u : ∂sol)
+    end
+    return primal, solve_up_adjoint
 end
 
 end

lib/NonlinearSolveBase/ext/NonlinearSolveBaseMooncakeExt.jl

@@ -2,30 +2,25 @@ module NonlinearSolveBaseMooncakeExt
 
 using NonlinearSolveBase, Mooncake
 using SciMLBase: SciMLBase
-import Mooncake: rrule!!, CoDual, zero_fcodual, @is_primitive,
-                 @from_rrule, @zero_adjoint, @mooncake_overlay, MinimalCtx,
-                 NoPullback
+using Mooncake: rrule!!, CoDual, zero_fcodual, @is_primitive,
+    @from_chainrules, @zero_adjoint, @mooncake_overlay, MinimalCtx,
+    NoPullback
 
-@from_rrule(MinimalCtx,
-    Tuple{
-        typeof(NonlinearSolveBase.solve_up),
-        SciMLBase.AbstractNonlinearProblem,
-        Union{Nothing, SciMLBase.AbstractSensitivityAlgorithm},
-        Any,
-        Any,
-        Any
-    },
-    true,)
+@from_chainrules MinimalCtx Tuple{typeof(NonlinearSolveBase.solve_up),
+    SciMLBase.AbstractNonlinearProblem,
+    Union{Nothing,SciMLBase.AbstractSensitivityAlgorithm},
+    Any,
+    Any,
+    Any
+} true
 
 # Dispatch for auto-alg
-@from_rrule(MinimalCtx,
-    Tuple{
-        typeof(NonlinearSolveBase.solve_up),
-        SciMLBase.AbstractNonlinearProblem,
-        Union{Nothing, SciMLBase.AbstractSensitivityAlgorithm},
-        Any,
-        Any
-    },
-    true,)
+@from_chainrules MinimalCtx Tuple{
+    typeof(NonlinearSolveBase.solve_up),
+    SciMLBase.AbstractNonlinearProblem,
+    Union{Nothing,SciMLBase.AbstractSensitivityAlgorithm},
+    Any,
+    Any
+} true
 
 end

test/adjoint_tests.jl

@@ -26,7 +26,7 @@
 
         @test ∂p_zygote ≈ ∂p_tracker ≈ ∂p_reversediff ≈ ∂p_enzyme
         @test ∂p_zygote ≈ ∂p_forwarddiff ≈ ∂p_tracker ≈ ∂p_reversediff ≈ ∂p_enzyme
-        @test_broken ∂p_forwarddiff ≈ ∂p_mooncake
+        @test ∂p_forwarddiff ≈ ∂p_mooncake
     else
         @info "Skipping adjoint tests on Julia $(VERSION) - Enzyme/SciMLSensitivity not compatible with 1.12+"
     end