Updating SciML to ChainRules

[ChrisRackauckas]

The list of things to handle is:

• https://github.com/SciML/DiffEqBase.jl/blob/44d1cf859c7073249a5455ce49841237f4693284/src/zygote.jl
• https://github.com/SciML/RecursiveArrayTools.jl/blob/734999e278ab5a6f9197669df155341d284b26fd/src/zygote.jl
• https://github.com/SciML/Quadrature.jl/blob/8309100db8e3575bf63d8f4ac22ab1f6cb15e43e/src/Quadrature.jl#L488-L563
• https://github.com/SciML/DiffEqGPU.jl/blob/8e3cb0d362f0413e57ea920bd6c216eee8cb97af/src/DiffEqGPU.jl#L344-L401
• https://github.com/SciML/DiffEqBase.jl/blob/master/src/solve.jl#L290-L323

The solve.jl one is a bit nasty because it calls all of the adjoints in DiffEqSensitivity.jl, so it's somewhat breaking, so we can... just do it really quickly.

@frankschae @YingboMa

[ChrisRackauckas]

@oxinabox for comfort.

[ChrisRackauckas]

SciML/DiffEqBase.jl#674
SciML/Integrals.jl#66

[ChrisRackauckas]

SciML/RecursiveArrayTools.jl#144
SciML/DiffEqGPU.jl#105

[ChrisRackauckas]

The current blocking issues:

• No translation for literal_getproperty. Keeping those to ZygoteRules.@adjoint for now
• https://github.com/SciML/DiffEqBase.jl/pull/674/files#diff-ad6ae73d4c409d0f5cdf63e58b8d5c3bbbe886a59c48aa5d0cbc3594a4f0caffR29-R31 translating map rules to not require __context__
• ZygoteRules.@adjoint function VectorOfArray(u) -> function ChainRulesCore.rrule(::VectorOfArray,u) isn't caught? https://github.com/SciML/RecursiveArrayTools.jl/pull/144/checks?check_run_id=2878240138#step:6:241

[ChrisRackauckas]

https://github.com/SciML/DiffEqGPU.jl/blob/master/src/DiffEqGPU.jl#L160-L162 doesn't get called, so https://github.com/SciML/DiffEqGPU.jl/blob/master/src/DiffEqGPU.jl#L164-L166 is required.

[ChrisRackauckas]

ZygoteRules.@adjoint function VectorOfArray(u) -> function ChainRulesCore.rrule(::Type{<:VectorOfArray},u)

Downstream testing revealed that Zygote actually ignored these rules 🤦 so those need to get reverted.

[ChrisRackauckas]

The end result of this is that in order to pass tests, all packages needed to keep a few ZygoteRules so it doesn't seem possible at this time to use strictly ChainRules for all of this. This should be double checked in the future (with Diffractor)

[ChrisRackauckas]

using OrdinaryDiffEq, DiffEqSensitivity, DiffEqFlux, LinearAlgebra, Flux
nn = FastChain(FastDense(1,16),FastDense(16,16,tanh),FastDense(16,2))
initial = initial_params(nn)

function ode2!(u, p, t)
    f1, f2 = nn([t],p)
    [-f1^2; f2]
end

tspan = (0.0, 10.0)
prob = ODEProblem(ode2!, Complex{Float64}[0;0], tspan, initial)

function loss(p)
  sol = last(solve(prob, Tsit5(), p=p, sensealg=BacksolveAdjoint(autojacvec=ZygoteVJP())))
  return norm(sol)
end

result_ode = DiffEqFlux.sciml_train(loss, initial, ADAM(0.1), maxiters = 100)

showed that Zygote skips

function ChainRulesCore.rrule(::typeof(getindex),VA::AbstractVectorOfArray, i)
  function AbstractVectorOfArray_getindex_adjoint(Δ)
    Δ′ = [ (i == j ? Δ : zero(x)) for (x,j) in zip(VA.u, 1:length(VA))]
    (NoTangent(),Δ′,NoTangent())
  end
  VA[i],AbstractVectorOfArray_getindex_adjoint
end

function ChainRulesCore.rrule(::typeof(getindex),VA::AbstractVectorOfArray, i, j...)
  function AbstractVectorOfArray_getindex_adjoint(Δ)
    Δ′ = zero(VA)
    Δ′[i,j...] = Δ
    (NoTangent(), Δ′, i,map(_ -> NoTangent(), j)...)
  end
  VA[i,j...],AbstractVectorOfArray_getindex_adjoint
end

because it has its own getindex overloads.

[ChrisRackauckas]

Zygote ignores:

function ChainRulesCore.rrule(::Type{<:EnsembleGPUArray})
    EnsembleGPUArray(0.0), _ -> NoTangent()
end

so I needed to also keep:

ZygoteRules.@adjoint function EnsembleGPUArray()
    EnsembleGPUArray(0.0), _ -> nothing
end

[ChrisRackauckas]

function ChainRulesCore.rrule(f::ODEFunction,u,p,t)
  if f.vjp === nothing
    ChainRulesCore.rrule(f.f,u,p,t)
  else
    f.vjp(u,p,t)
  end
end

was skipped, so I added back:

ZygoteRules.@adjoint function (f::ODEFunction)(u,p,t)
  if f.vjp === nothing
    ZygoteRules._pullback(f.f,u,p,t)
  else
    f.vjp(u,p,t)
  end
end

ZygoteRules.@adjoint! function (f::ODEFunction)(du,u,p,t)
  if f.vjp === nothing
    ZygoteRules._pullback(f.f,du,u,p,t)
  else
    f.vjp(du,u,p,t)
  end
end