fix adjoints

Author: jClugstor

lib/NonlinearSolveBase/src/NonlinearSolveBase.jl

@@ -23,7 +23,7 @@ using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinear
                  LinearAliasSpecifier, ImmutableNonlinearProblem, NonlinearAliasSpecifier,
                  promote_u0, get_concrete_u0, get_concrete_p,
                  has_kwargs, extract_alg, promote_u0, checkkwargs, SteadyStateProblem,
-                 NoDefaultAlgorithmError, NonSolverError, KeywordArgError
+                 NoDefaultAlgorithmError, NonSolverError, KeywordArgError, AbstractDEAlgorithm
 import SciMLBase: solve, init, __init, __solve, wrap_sol, get_root_indp, isinplace, remake
 
 using SciMLJacobianOperators: JacobianOperator, StatefulJacobianOperator

lib/NonlinearSolveBase/src/solve.jl

@@ -794,11 +794,11 @@ end
 function _solve_adjoint(prob, sensealg, u0, p, originator, args...; merge_callbacks = true,
         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,
+    if isnothing(alg) || !(alg isa AbstractNonlinearAlgorithm) # Default algorithm handling
+        _prob = get_concrete_problem(prob, true; u0 = u0,
             p = p, kwargs...)
     else
-        _prob = get_concrete_problem(prob, isadaptive(alg); u0 = u0, p = p, kwargs...)
+        _prob = get_concrete_problem(prob, true; u0 = u0, p = p, kwargs...)
     end
 
     if has_kwargs(_prob)

lib/NonlinearSolveBase/src/termination_conditions.jl

@@ -82,7 +82,7 @@ function CommonSolve.init(
     length(saved_value_prototype) == 0 && (saved_value_prototype = nothing)
 
     leastsq = typeof(prob) <: NonlinearLeastSquaresProblem
-
+    Main.@infiltrate
     return NonlinearTerminationModeCache(
         u_unaliased, ReturnCode.Default, abstol, reltol, best_value, mode,
         initial_objective, objectives_trace, 0, saved_value_prototype,

lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveChainRulesCoreExt.jl

@@ -2,19 +2,18 @@ module SimpleNonlinearSolveChainRulesCoreExt
 
 using ChainRulesCore: ChainRulesCore, NoTangent
 
-using NonlinearSolveBase: ImmutableNonlinearProblem
+using NonlinearSolveBase: ImmutableNonlinearProblem, _solve_adjoint
 using SciMLBase: ChainRulesOriginator, NonlinearLeastSquaresProblem
 
-using SimpleNonlinearSolve: SimpleNonlinearSolve, simplenonlinearsolve_solve_up,
-                            solve_adjoint
+using SimpleNonlinearSolve: SimpleNonlinearSolve, simplenonlinearsolve_solve_up
 
 function ChainRulesCore.rrule(
         ::typeof(simplenonlinearsolve_solve_up),
         prob::Union{ImmutableNonlinearProblem, NonlinearLeastSquaresProblem},
         sensealg, u0, u0_changed, p, p_changed, alg, args...; kwargs...
 )
     out,
-    ∇internal = solve_adjoint(
+    ∇internal = _solve_adjoint(
         prob, sensealg, u0, p, ChainRulesOriginator(), alg, args...; kwargs...
     )
     function ∇simplenonlinearsolve_solve_up(Δ)

lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveDiffEqBaseExt.jl

@@ -2,12 +2,12 @@ module SimpleNonlinearSolveDiffEqBaseExt
 
 #using DiffEqBase: DiffEqBase
 
-using SimpleNonlinearSolve: SimpleNonlinearSolve
+# using SimpleNonlinearSolve: SimpleNonlinearSolve
 
-SimpleNonlinearSolve.is_extension_loaded(::Val{:DiffEqBase}) = true
+# SimpleNonlinearSolve.is_extension_loaded(::Val{:DiffEqBase}) = true
 
-function SimpleNonlinearSolve.solve_adjoint_internal(args...; kwargs...)
-    return DiffEqBase._solve_adjoint(args...; kwargs...)
-end
+# function SimpleNonlinearSolve.solve_adjoint_internal(args...; kwargs...)
+#     return DiffEqBase._solve_adjoint(args...; kwargs...)
+# end
 
 end

lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveReverseDiffExt.jl

@@ -1,12 +1,12 @@
 module SimpleNonlinearSolveReverseDiffExt
 
-using NonlinearSolveBase: ImmutableNonlinearProblem
+using NonlinearSolveBase: ImmutableNonlinearProblem, _solve_adjoint
 using SciMLBase: ReverseDiffOriginator, NonlinearLeastSquaresProblem, remake
 
 using ArrayInterface: ArrayInterface
 using ReverseDiff: ReverseDiff, TrackedArray, TrackedReal
 
-using SimpleNonlinearSolve: SimpleNonlinearSolve, solve_adjoint
+using SimpleNonlinearSolve: SimpleNonlinearSolve
 import SimpleNonlinearSolve: simplenonlinearsolve_solve_up
 
 for pType in (ImmutableNonlinearProblem, NonlinearLeastSquaresProblem)
@@ -27,7 +27,7 @@ for pType in (ImmutableNonlinearProblem, NonlinearLeastSquaresProblem)
         u0, p = ReverseDiff.value(tu0), ReverseDiff.value(tp)
         prob = remake(tprob; u0, p)
         out,
-        ∇internal = solve_adjoint(
+        ∇internal = _solve_adjoint(
             prob, sensealg, u0, p, ReverseDiffOriginator(), alg, args...; kwargs...)
 
         function ∇simplenonlinearsolve_solve_up(Δ...)

lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveTrackerExt.jl

@@ -1,12 +1,12 @@
 module SimpleNonlinearSolveTrackerExt
 
-using NonlinearSolveBase: ImmutableNonlinearProblem
+using NonlinearSolveBase: ImmutableNonlinearProblem, _solve_adjoint
 using SciMLBase: TrackerOriginator, NonlinearLeastSquaresProblem, remake
 
 using ArrayInterface: ArrayInterface
 using Tracker: Tracker, TrackedArray, TrackedReal
 
-using SimpleNonlinearSolve: SimpleNonlinearSolve, solve_adjoint
+using SimpleNonlinearSolve: SimpleNonlinearSolve
 
 for pType in (ImmutableNonlinearProblem, NonlinearLeastSquaresProblem)
     aTypes = (TrackedArray, AbstractArray{<:TrackedReal}, Any)
@@ -26,7 +26,7 @@ for pType in (ImmutableNonlinearProblem, NonlinearLeastSquaresProblem)
         u0, p = Tracker.data(tu0), Tracker.data(tp)
         prob = remake(tprob; u0, p)
         out,
-        ∇internal = solve_adjoint(
+        ∇internal = _solve_adjoint(
             prob, sensealg, u0, p, TrackerOriginator(), alg, args...; kwargs...)
 
         function ∇simplenonlinearsolve_solve_up(Δ)

lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl

@@ -125,12 +125,12 @@ end
 
 # NOTE: This is defined like this so that we don't have to keep have 2 args for the
 # extensions
-function solve_adjoint(args...; kws...)
-    is_extension_loaded(Val(:DiffEqBase)) && return solve_adjoint_internal(args...; kws...)
-    error("Adjoint sensitivity analysis requires `DiffEqBase.jl` to be explicitly loaded.")
-end
+# function solve_adjoint(args...; kws...)
+#     is_extension_loaded(Val(:DiffEqBase)) && return solve_adjoint_internal(args...; kws...)
+#     error("Adjoint sensitivity analysis requires `DiffEqBase.jl` to be explicitly loaded.")
+# end
 
-function solve_adjoint_internal end
+# function solve_adjoint_internal end
 
 @setup_workload begin
     for T in (Float64,)

lib/SimpleNonlinearSolve/test/core/adjoint_tests.jl

@@ -1,5 +1,5 @@
 @testitem "Simple Adjoint Test" tags=[:adjoint] begin
-    using ForwardDiff, ReverseDiff, SciMLSensitivity, Tracker, Zygote, DiffEqBase
+    using ForwardDiff, ReverseDiff, SciMLSensitivity, Tracker, Zygote
 
     ff(u, p) = u .^ 2 .- p
 

test/adjoint_tests.jl

@@ -0,0 +1,19 @@
+using ForwardDiff, ReverseDiff, SciMLSensitivity, Tracker, Zygote, NonlinearSolve, Test
+
+ff(u, p) = u .^ 2 .- p
+
+function solve_nlprob(p)
+    prob = NonlinearProblem{false}(ff, [1.0, 2.0], p)
+    sol = solve(prob, NewtonRaphson())
+    res = sol isa AbstractArray ? sol : sol.u
+    return sum(abs2, res)
+end
+
+p = [3.0, 2.0]
+
+∂p_zygote = only(Zygote.gradient(solve_nlprob, p))
+∂p_forwarddiff = ForwardDiff.gradient(solve_nlprob, p)
+∂p_tracker = Tracker.data(only(Tracker.gradient(solve_nlprob, p)))
+∂p_reversediff = ReverseDiff.gradient(solve_nlprob, p)
+@test ∂p_zygote ≈ ∂p_tracker ≈ ∂p_reversediff
+@test ∂p_zygote ≈ ∂p_forwarddiff ≈ ∂p_tracker ≈ ∂p_reversediff