[WIP] Tracker and ReverseDiff custom concrete_solve gradients

MWE:

```julia
using DiffEqSensitivity, OrdinaryDiffEq, Zygote
using RecursiveArrayTools: DiffEqArray
using Test, ForwardDiff
using Tracker, ReverseDiff

function fiip(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
end
function foop(u,p,t)
  dx = p[1]*u[1] - p[2]*u[1]*u[2]
  dy = -p[3]*u[2] + p[4]*u[1]*u[2]
  [dx,dy]
end
function foop(u::Tracker.TrackedArray,p,t)
  dx = p[1]*u[1] - p[2]*u[1]*u[2]
  dy = -p[3]*u[2] + p[4]*u[1]*u[2]
  Tracker.collect([dx,dy])
end

p = [1.5,1.0,3.0,1.0]; u0 = [1.0;1.0]
prob = ODEProblem(fiip,u0,(0.0,10.0),p)
proboop = ODEProblem(foop,u0,(0.0,10.0),p)

sol = concrete_solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)

_sol = solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14)
ū0,adj = adjoint_sensitivities(_sol,Tsit5(),((out,u,p,t,i) -> out .= -1),0.0:0.1:10,abstol=1e-14,
                                  reltol=1e-14,iabstol=1e-14,ireltol=1e-12)

du01,dp1 = Zygote.gradient((u0,p)->sum(concrete_solve(prob,Tsit5(),u0,p,abstol=1e-14,reltol=1e-14,saveat=0.1)),u0,p)

Tracker.@Grad function concrete_solve(prob,alg,u0,p,args...;
                                      sensealg=nothing,kwargs...)
  @show "here"
  _concrete_solve_adjoint(prob,alg,sensealg,u0,p,args...;kwargs...)
end

ReverseDiff.@Grad function concrete_solve(prob,alg,u0,p,args...;
                                          sensealg=nothing,kwargs...)
  @show "here"
  _concrete_solve_adjoint(prob,alg,sensealg,u0,p,args...;kwargs...)
end

du01,dp1 = Tracker.gradient((u0,p)->sum(concrete_solve(prob,Tsit5(),u0,p,abstol=1e-14,reltol=1e-14,saveat=0.1)),u0,p)
du01,dp1 = ReverseDiff.gradient((u0,p)->sum(concrete_solve(prob,Tsit5(),u0,p,abstol=1e-14,reltol=1e-14,saveat=0.1)),(u0,p))

Tracker.@Grad function concrete_solve(prob,alg,u0,p)
  @show "here"
  _concrete_solve_adjoint(prob,alg,sensealg,u0,p)
end

ReverseDiff.@Grad function concrete_solve(prob,alg,u0,p)
  @show "here"
  _concrete_solve_adjoint(prob,alg,sensealg,u0,p)
end

du01,dp1 = Tracker.gradient((u0,p)->sum(concrete_solve(prob,Tsit5(),u0,p)),u0,p)
du01,dp1 = ReverseDiff.gradient((u0,p)->sum(concrete_solve(prob,Tsit5(),u0,p)),(u0,p))
```

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "DiffEqBase"
 uuid = "2b5f629d-d688-5b77-993f-72d75c75574e"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "6.34.3"
+version = "6.35.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/init.jl

@@ -145,37 +145,7 @@ function __init__()
   end
 
   @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" begin
-    value(x::Type{Tracker.TrackedReal{T}}) where T = T
-    value(x::Type{Tracker.TrackedArray{T,N,A}}) where {T,N,A} = Array{T,N}
-    value(x::Tracker.TrackedReal)  = x.data
-    value(x::Tracker.TrackedArray) = x.data
-
-    @inline fastpow(x::Tracker.TrackedReal, y::Tracker.TrackedReal) = x^y
-    @inline Base.any(f::Function,x::Tracker.TrackedArray) = any(f,Tracker.data(x))
-
-    # Support adaptive with non-tracked time
-    @inline function ODE_DEFAULT_NORM(u::Tracker.TrackedArray,t) where {N}
-      sqrt(sum(abs2,value(u)) / length(u))
-    end
-    @inline function ODE_DEFAULT_NORM(u::AbstractArray{<:Tracker.TrackedReal,N},t) where {N}
-      sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip((value(x) for x in u),Iterators.repeated(t))) / length(u))
-    end
-    @inline function ODE_DEFAULT_NORM(u::Array{<:Tracker.TrackedReal,N},t) where {N}
-      sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip((value(x) for x in u),Iterators.repeated(t))) / length(u))
-    end
-    @inline ODE_DEFAULT_NORM(u::Tracker.TrackedReal,t) = abs(value(u))
-
-    # Support TrackedReal time, don't drop tracking on the adaptivity there
-    @inline function ODE_DEFAULT_NORM(u::Tracker.TrackedArray,t::Tracker.TrackedReal) where {N}
-      sqrt(sum(abs2,u) / length(u))
-    end
-    @inline function ODE_DEFAULT_NORM(u::AbstractArray{<:Tracker.TrackedReal,N},t::Tracker.TrackedReal) where {N}
-      sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip(u,Iterators.repeated(t))) / length(u))
-    end
-    @inline function ODE_DEFAULT_NORM(u::Array{<:Tracker.TrackedReal,N},t::Tracker.TrackedReal) where {N}
-      sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip(u,Iterators.repeated(t))) / length(u))
-    end
-    @inline ODE_DEFAULT_NORM(u::Tracker.TrackedReal,t::Tracker.TrackedReal) = abs(u)
+    include("tracker.jl")
   end
 
   # Piracy, should get upstreamed
@@ -213,6 +183,10 @@ function __init__()
     end
   end
 
+  @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
+    include("reversediff.jl")
+  end
+
   @require GeneralizedGenerated="6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb" begin
     numargs(::GeneralizedGenerated.RuntimeFn{Args}) where Args = GeneralizedGenerated.from_type(Args) |> length
   end

src/reversediff.jl

@@ -0,0 +1,4 @@
+ReverseDiff.@grad function concrete_solve(prob,alg,u0,p,args...;
+                                          sensealg=nothing,kwargs...)
+  _concrete_solve_adjoint(prob,alg,sensealg,u0,p,args...;kwargs...)
+end

src/tracker.jl

@@ -0,0 +1,36 @@
+value(x::Type{Tracker.TrackedReal{T}}) where T = T
+value(x::Type{Tracker.TrackedArray{T,N,A}}) where {T,N,A} = Array{T,N}
+value(x::Tracker.TrackedReal)  = x.data
+value(x::Tracker.TrackedArray) = x.data
+
+@inline fastpow(x::Tracker.TrackedReal, y::Tracker.TrackedReal) = x^y
+@inline Base.any(f::Function,x::Tracker.TrackedArray) = any(f,Tracker.data(x))
+
+# Support adaptive with non-tracked time
+@inline function ODE_DEFAULT_NORM(u::Tracker.TrackedArray,t) where {N}
+  sqrt(sum(abs2,value(u)) / length(u))
+end
+@inline function ODE_DEFAULT_NORM(u::AbstractArray{<:Tracker.TrackedReal,N},t) where {N}
+  sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip((value(x) for x in u),Iterators.repeated(t))) / length(u))
+end
+@inline function ODE_DEFAULT_NORM(u::Array{<:Tracker.TrackedReal,N},t) where {N}
+  sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip((value(x) for x in u),Iterators.repeated(t))) / length(u))
+end
+@inline ODE_DEFAULT_NORM(u::Tracker.TrackedReal,t) = abs(value(u))
+
+# Support TrackedReal time, don't drop tracking on the adaptivity there
+@inline function ODE_DEFAULT_NORM(u::Tracker.TrackedArray,t::Tracker.TrackedReal) where {N}
+  sqrt(sum(abs2,u) / length(u))
+end
+@inline function ODE_DEFAULT_NORM(u::AbstractArray{<:Tracker.TrackedReal,N},t::Tracker.TrackedReal) where {N}
+  sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip(u,Iterators.repeated(t))) / length(u))
+end
+@inline function ODE_DEFAULT_NORM(u::Array{<:Tracker.TrackedReal,N},t::Tracker.TrackedReal) where {N}
+  sqrt(sum(x->ODE_DEFAULT_NORM(x[1],x[2]),zip(u,Iterators.repeated(t))) / length(u))
+end
+@inline ODE_DEFAULT_NORM(u::Tracker.TrackedReal,t::Tracker.TrackedReal) = abs(u)
+
+Tracker.@grad function concrete_solve(prob,alg,u0,p,args...;
+                                      sensealg=nothing,kwargs...)
+  _concrete_solve_adjoint(prob,alg,sensealg,u0,p,args...;kwargs...)
+end