add Mooncake to adjoint tests

Author: jClugstor

Project.toml

@@ -99,6 +99,7 @@ LineSearches = "7.3"
 LinearAlgebra = "1.10"
 LinearSolve = "2.36.1, 3"
 MINPACK = "1.2"
+Mooncake = "0.4"
 MPI = "0.20.22"
 NLSolvers = "0.5"
 NLsolve = "4.5"
@@ -146,6 +147,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 NLSolvers = "337daf1e-9722-11e9-073e-8b9effe078ba"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
@@ -170,4 +172,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SparseMatrixColorings", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote", "ReverseDiff", "Tracker", "SciMLSensitivity", "Enzyme"]
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SparseMatrixColorings", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote", "ReverseDiff", "Tracker", "SciMLSensitivity", "Enzyme", "Mooncake"]

lib/NonlinearSolveBase/ext/NonlinearSolveBaseMooncakeExt.jl

@@ -1,7 +1,7 @@
 module NonlinearSolveBaseMooncakeExt
 
 using NonlinearSolveBase, Mooncake
-using SciMLBase: ADOriginator, ChainRulesOriginator, MooncakeOriginator
+using SciMLBase: SciMLBase
 import Mooncake: rrule!!, CoDual, zero_fcodual, @is_primitive,
                  @from_rrule, @zero_adjoint, @mooncake_overlay, MinimalCtx,
                  NoPullback

lib/NonlinearSolveBase/src/solve.jl

@@ -795,7 +795,7 @@ function _solve_adjoint(prob, sensealg, u0, p, originator, args...; merge_callba
         kwargs...)
     alg = extract_alg(args, kwargs, prob.kwargs)
     if isnothing(alg) || !(alg isa AbstractDEAlgorithm) # Default algorithm handling
-        _prob = get_concrete_problem(prob, !(prob isa DiscreteProblem); u0 = u0,
+        _prob = get_concrete_problem(prob, true; u0 = u0,
             p = p, kwargs...)
     else
         _prob = get_concrete_problem(prob, isadaptive(alg); u0 = u0, p = p, kwargs...)
@@ -817,7 +817,7 @@ function _solve_forward(prob, sensealg, u0, p, originator, args...; merge_callba
         kwargs...)
     alg = extract_alg(args, kwargs, prob.kwargs)
     if isnothing(alg) || !(alg isa AbstractDEAlgorithm) # Default algorithm handling
-        _prob = get_concrete_problem(prob, !(prob isa DiscreteProblem); u0 = u0,
+        _prob = get_concrete_problem(prob, true; u0 = u0,
             p = p, kwargs...)
     else
         _prob = get_concrete_problem(prob, isadaptive(alg); u0 = u0, p = p, kwargs...)

test/adjoint_tests.jl

@@ -1,5 +1,5 @@
-@testitem "Adjoint Tests" tags = [:adjoint] begin
-    using ForwardDiff, ReverseDiff, SciMLSensitivity, Tracker, Zygote, Enzyme
+@testitem "Adjoint Tests" tags = [:nopre] begin
+    using ForwardDiff, ReverseDiff, SciMLSensitivity, Tracker, Zygote, Enzyme, Mooncake
 
     ff(u, p) = u .^ 2 .- p
 
@@ -17,6 +17,11 @@
     ∂p_tracker = Tracker.data(only(Tracker.gradient(solve_nlprob, p)))
     ∂p_reversediff = ReverseDiff.gradient(solve_nlprob, p)
     ∂p_enzyme = Enzyme.gradient(Enzyme.Reverse, solve_nlprob, p)[1]
-    @test ∂p_zygote ≈ ∂p_tracker ≈ ∂p_reversediff ≈ ∂p_enzyme
+
+    cache = Mooncake.prepare_gradient_cache(solve_nlprob, p)
+    ∂p_mooncake = Mooncake.value_and_gradient!!(cache, solve_nlprob, p)[2][2]
+
+    @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
 end