Merge pull request #133 from SciML/ap/reversediff

Automatic Differentiation

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.5.0"
+version = "1.6.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -22,20 +22,24 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore"
 SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
+SimpleNonlinearSolveReverseDiffExt = "ReverseDiff"
 SimpleNonlinearSolveStaticArraysExt = "StaticArrays"
+SimpleNonlinearSolveTrackerExt = "Tracker"
 SimpleNonlinearSolveZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "0.2.6"
 AllocCheck = "0.1.1"
-ArrayInterface = "7.7"
 Aqua = "0.8"
+ArrayInterface = "7.7"
 CUDA = "5.2"
 ChainRulesCore = "1.22"
 ConcreteStructs = "0.2.3"
@@ -54,11 +58,13 @@ PrecompileTools = "1.2"
 Random = "1.10"
 ReTestItems = "1.23"
 Reexport = "1.2"
+ReverseDiff = "1.15"
 SciMLBase = "2.26.3"
 SciMLSensitivity = "7.56"
 StaticArrays = "1.9"
 StaticArraysCore = "1.4.2"
 Test = "1.10"
+Tracker = "0.2.32"
 Zygote = "0.6.69"
 julia = "1.10"
 
@@ -77,10 +83,12 @@ PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "AllocCheck", "DiffEqBase", "ForwardDiff", "LinearAlgebra", "LinearSolve", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SciMLSensitivity", "StaticArrays", "Zygote", "CUDA", "PolyesterForwardDiff", "Reexport", "Test", "FiniteDiff"]
+test = ["Aqua", "AllocCheck", "DiffEqBase", "ForwardDiff", "LinearAlgebra", "LinearSolve", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SciMLSensitivity", "StaticArrays", "Zygote", "CUDA", "PolyesterForwardDiff", "Reexport", "Test", "FiniteDiff", "ReverseDiff", "Tracker"]

ext/SimpleNonlinearSolveChainRulesCoreExt.jl

@@ -5,14 +5,13 @@ using ChainRulesCore, DiffEqBase, SciMLBase, SimpleNonlinearSolve
 # The expectation here is that no-one is using this directly inside a GPU kernel. We can
 # eventually lift this requirement using a custom adjoint
 function ChainRulesCore.rrule(::typeof(SimpleNonlinearSolve.__internal_solve_up),
-        prob::NonlinearProblem,
-        sensealg::Union{Nothing, DiffEqBase.AbstractSensitivityAlgorithm}, u0, u0_changed,
-        p, p_changed, alg, args...; kwargs...)
+        prob::NonlinearProblem, sensealg, u0, u0_changed, p, p_changed, alg, args...;
+        kwargs...)
     out, ∇internal = DiffEqBase._solve_adjoint(prob, sensealg, u0, p,
         SciMLBase.ChainRulesOriginator(), alg, args...; kwargs...)
     function ∇__internal_solve_up(Δ)
         ∂f, ∂prob, ∂sensealg, ∂u0, ∂p, ∂originator, ∂args... = ∇internal(Δ)
-        return (∂f, ∂prob, ∂sensealg, ∂u0, NoTangent(), ∂p, NoTangent(), ∂originator,
+        return (∂f, ∂prob, ∂sensealg, ∂u0, NoTangent(), ∂p, NoTangent(), NoTangent(),
             ∂args...)
     end
     return out, ∇__internal_solve_up

ext/SimpleNonlinearSolveReverseDiffExt.jl

@@ -0,0 +1,60 @@
+module SimpleNonlinearSolveReverseDiffExt
+
+using ArrayInterface, DiffEqBase, ReverseDiff, SciMLBase, SimpleNonlinearSolve
+import ReverseDiff: TrackedArray, TrackedReal
+import SimpleNonlinearSolve: __internal_solve_up
+
+function __internal_solve_up(
+        prob::NonlinearProblem, sensealg, u0::TrackedArray, u0_changed,
+        p::TrackedArray, p_changed, alg, args...; kwargs...)
+    return ReverseDiff.track(__internal_solve_up, prob, sensealg, u0,
+        u0_changed, p, p_changed, alg, args...; kwargs...)
+end
+
+function __internal_solve_up(
+        prob::NonlinearProblem, sensealg, u0, u0_changed,
+        p::TrackedArray, p_changed, alg, args...; kwargs...)
+    return ReverseDiff.track(__internal_solve_up, prob, sensealg, u0,
+        u0_changed, p, p_changed, alg, args...; kwargs...)
+end
+
+function __internal_solve_up(
+        prob::NonlinearProblem, sensealg, u0::TrackedArray, u0_changed,
+        p, p_changed, alg, args...; kwargs...)
+    return ReverseDiff.track(__internal_solve_up, prob, sensealg, u0,
+        u0_changed, p, p_changed, alg, args...; kwargs...)
+end
+
+function __internal_solve_up(prob::NonlinearProblem, sensealg,
+        u0::AbstractArray{<:TrackedReal}, u0_changed, p::AbstractArray{<:TrackedReal},
+        p_changed, alg, args...; kwargs...)
+    return __internal_solve_up(
+        prob, sensealg, ArrayInterface.aos_to_soa(u0), true,
+        ArrayInterface.aos_to_soa(p), true, alg, args...; kwargs...)
+end
+
+function __internal_solve_up(prob::NonlinearProblem, sensealg, u0, u0_changed,
+        p::AbstractArray{<:TrackedReal}, p_changed, alg, args...; kwargs...)
+    return __internal_solve_up(
+        prob, sensealg, u0, true, ArrayInterface.aos_to_soa(p), true, alg, args...; kwargs...)
+end
+
+function __internal_solve_up(prob::NonlinearProblem, sensealg,
+        u0::AbstractArray{<:TrackedReal}, u0_changed, p, p_changed, alg, args...; kwargs...)
+    return __internal_solve_up(
+        prob, sensealg, u0, true, ArrayInterface.aos_to_soa(p), true, alg, args...; kwargs...)
+end
+
+ReverseDiff.@grad function __internal_solve_up(
+        prob::NonlinearProblem, sensealg, u0, u0_changed, p, p_changed, alg, args...; kwargs...)
+    out, ∇internal = DiffEqBase._solve_adjoint(
+        prob, sensealg, ReverseDiff.value(u0), ReverseDiff.value(p),
+        SciMLBase.ReverseDiffOriginator(), alg, args...; kwargs...)
+    function ∇__internal_solve_up(_args...)
+        ∂prob, ∂sensealg, ∂u0, ∂p, ∂originator, ∂args... = ∇internal(_args...)
+        return (∂prob, ∂sensealg, ∂u0, nothing, ∂p, nothing, nothing, ∂args...)
+    end
+    return Array(out), ∇__internal_solve_up
+end
+
+end

ext/SimpleNonlinearSolveTrackerExt.jl

@@ -0,0 +1,42 @@
+module SimpleNonlinearSolveTrackerExt
+
+using DiffEqBase, SciMLBase, SimpleNonlinearSolve, Tracker
+
+function SimpleNonlinearSolve.__internal_solve_up(prob::NonlinearProblem,
+        sensealg, u0::TrackedArray, u0_changed, p, p_changed, alg, args...; kwargs...)
+    return Tracker.track(
+        SimpleNonlinearSolve.__internal_solve_up, prob, sensealg, u0, u0_changed,
+        p, p_changed, alg, args...; kwargs...)
+end
+
+function SimpleNonlinearSolve.__internal_solve_up(
+        prob::NonlinearProblem, sensealg, u0::TrackedArray, u0_changed,
+        p::TrackedArray, p_changed, alg, args...; kwargs...)
+    return Tracker.track(
+        SimpleNonlinearSolve.__internal_solve_up, prob, sensealg, u0, u0_changed,
+        p, p_changed, alg, args...; kwargs...)
+end
+
+function SimpleNonlinearSolve.__internal_solve_up(prob::NonlinearProblem,
+        sensealg, u0, u0_changed, p::TrackedArray, p_changed, alg, args...; kwargs...)
+    return Tracker.track(
+        SimpleNonlinearSolve.__internal_solve_up, prob, sensealg, u0, u0_changed,
+        p, p_changed, alg, args...; kwargs...)
+end
+
+Tracker.@grad function SimpleNonlinearSolve.__internal_solve_up(_prob::NonlinearProblem,
+        sensealg, u0_, u0_changed, p_, p_changed, alg, args...; kwargs...)
+    u0, p = Tracker.data(u0_), Tracker.data(p_)
+    prob = remake(_prob; u0, p)
+    out, ∇internal = DiffEqBase._solve_adjoint(prob, sensealg, u0, p,
+        SciMLBase.TrackerOriginator(), alg, args...; kwargs...)
+
+    function ∇__internal_solve_up(Δ)
+        ∂prob, ∂sensealg, ∂u0, ∂p, ∂originator, ∂args... = ∇internal(Δ)
+        return (∂prob, ∂sensealg, ∂u0, nothing, ∂p, nothing, nothing, ∂args...)
+    end
+
+    return out, ∇__internal_solve_up
+end
+
+end

src/SimpleNonlinearSolve.jl

@@ -71,8 +71,8 @@ function SciMLBase.solve(
         alg, args...; prob.kwargs..., kwargs...)
 end
 
-function __internal_solve_up(_prob::NonlinearProblem, sensealg, u0, u0_changed, p,
-        p_changed, alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...)
+function __internal_solve_up(_prob::NonlinearProblem, sensealg, u0, u0_changed,
+        p, p_changed, alg, args...; kwargs...)
     prob = u0_changed || p_changed ? remake(_prob; u0, p) : _prob
     return SciMLBase.__solve(prob, alg, args...; kwargs...)
 end

src/nlsolve/halley.jl

@@ -33,7 +33,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...;
     fx = _get_fx(prob, x)
     T = eltype(x)
 
-    autodiff = __get_concrete_autodiff(prob, alg.autodiff; polyester = Val(false))
+    autodiff = __get_concrete_autodiff(prob, alg.autodiff)
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fx, x,
         termination_condition)
 

src/utils.jl

@@ -365,18 +365,9 @@ end
 
 # Decide which AD backend to use
 @inline __get_concrete_autodiff(prob, ad::ADTypes.AbstractADType; kwargs...) = ad
-@inline function __get_concrete_autodiff(prob, ::Nothing; polyester::Val{P} = Val(true),
-        kwargs...) where {P}
-    if ForwardDiff.can_dual(eltype(prob.u0))
-        if P && __is_extension_loaded(Val(:PolyesterForwardDiff)) &&
-           !(prob.u0 isa Number) && ArrayInterface.can_setindex(prob.u0)
-            return AutoPolyesterForwardDiff()
-        else
-            return AutoForwardDiff()
-        end
-    else
-        return AutoFiniteDiff()
-    end
+@inline function __get_concrete_autodiff(prob, ::Nothing; kwargs...)
+    return ifelse(
+        ForwardDiff.can_dual(eltype(prob.u0)), AutoForwardDiff(), AutoFiniteDiff())
 end
 
 @inline __reshape(x::Number, args...) = x

test/core/adjoint_tests.jl

@@ -1,14 +1,21 @@
 @testitem "Simple Adjoint Test" begin
-    using ForwardDiff, SciMLSensitivity, Zygote
+    using ForwardDiff, ReverseDiff, SciMLSensitivity, Tracker, Zygote
 
     ff(u, p) = u .^ 2 .- p
 
     function solve_nlprob(p)
         prob = NonlinearProblem{false}(ff, [1.0, 2.0], p)
-        return sum(abs2, solve(prob, SimpleNewtonRaphson()).u)
+        sol = solve(prob, SimpleNewtonRaphson())
+        res = sol isa AbstractArray ? sol : sol.u
+        return sum(abs2, res)
     end
 
     p = [3.0, 2.0]
 
-    @test only(Zygote.gradient(solve_nlprob, p)) ≈ ForwardDiff.gradient(solve_nlprob, p)
+    ∂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_forwarddiff ≈ ∂p_tracker ≈ ∂p_reversediff
 end