ForwardDiff Overload Fixes

[jClugstor]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[jClugstor]

The overloads won't work with higher order derivatives or nested Duals. We would need a special way to handle higher order derivatives / maybe a general expression for x_p for higher order derivatives. So for now I just want to exclude nested Duals from going through this.
So I did this to try to exclude NestedDualLinearProblem from going through the overloads, but this is pretty restrictive. But if I do SingleDual{T, V, P} = Dual{T, V, P} where {T, V <:Number , P} then a NestedDualLinearProblem is also a NestedDualLinearProblem, which defeats the purpose. Essentially I want to make sure that if the value of the Dual is another Dual, go through the normal path and not the ForwardDiffExt path. @oscardssmith I wonder if you have any thoughts on how I could do that?

# Define type for non-nested dual numbers
const SingleDual{T, V, P} = Dual{T, V, P} where {T, V <:Float64 , P}

# Define type for nested dual numbers
const NestedDual{T, V, P} = Dual{T, V, P} where {T, V <: Dual, P}

const SingleDualLinearProblem = LinearProblem{
    <:Union{Number, <:AbstractArray, Nothing}, iip,
    <:Union{<:SingleDual, <:AbstractArray{<:SingleDual}},
    <:Union{<:SingleDual, <:AbstractArray{<:SingleDual}},
    <:Any
} where {iip}

const NestedDualLinearProblem = LinearProblem{
    <:Union{Number, <:AbstractArray, Nothing}, iip,
    <:Union{<:NestedDual, <:AbstractArray{<:NestedDual}},
    <:Union{<:NestedDual, <:AbstractArray{<:NestedDual}},
    <:Any
} where {iip}

const DualALinearProblem = LinearProblem{
    <:Union{Number, <:AbstractArray, Nothing},
    iip,
    <:Union{<:SingleDual, <:AbstractArray{<:SingleDual}},
    <:Union{Number, <:AbstractArray},
    <:Any
} where {iip}

const DualBLinearProblem = LinearProblem{
    <:Union{Number, <:AbstractArray, Nothing},
    iip,
    <:Union{Number, <:AbstractArray},
    <:Union{<:SingleDual, <:AbstractArray{<:SingleDual}},
    <:Any
} where {iip}

const DualAbstractLinearProblem = Union{
    SingleDualLinearProblem, DualALinearProblem, DualBLinearProblem}

[oscardssmith]

I think making the higher order/nested cases work out shouldn't be too hard. Derivatives are a linear operator, so I think everything should compose nicely here.

[jClugstor]

Using the general Leibniz rule for products I can get this:

For the partials we might need to use this:

I think handling the partials will be the trickiest part of this.

[jClugstor]

@oscardssmith this actually does work recursively, I just had to make a few adjustments first. I'll test it with more nested Dual stuff but this should be good pretty soon actually.

[jClugstor]

@oscardssmith this is mostly ready to go, but there is one problem with one of the SciMLSensitivity tests that doesn't run in the integration tests. https://github.com/SciML/SciMLSensitivity.jl/blob/e2916ef299f5e39998e2b758fd43e29e993a18c7/test/stiff_adjoints.jl#L24
In the stiff adjoint tests when I try to differentiate loss_function the solve becomes unstable: Warning: At t=1.5507202241421e-7, dt was forced below floating point epsilon 2.6469779601696886e-23 ... .

This only happens on this branch of LinearSolve, and so far it only happens with RadauIIA5 and RadauIIA9. RadauIIA3 works, and so do all of the other implicit solvers I've tested.

using SciMLSensitivity
using OrdinaryDiffEq, ForwardDiff, Test
using ForwardDiff: Dual
using Zygote


function lotka_volterra(u, p, t)
    x, y = u
    α, β, δ, γ = p
    [α * x - β * x * y, -δ * y + γ * x * y]
end

function lotka_volterra(du, u, p, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = α * x - β * x * y
    du[2] = -δ * y + γ * x * y
end

u0 = [1.0, 1.0];
tspan = (0.0, 10.0);
p0 = [1.5, 1.0, 3.0, 1.0];
prob0 = ODEProblem{true, SciMLBase.FullSpecialize}(lotka_volterra, u0, tspan, p0);
# Solve the ODE and collect solutions at fixed intervals
target_data = solve(prob0, FBDF(), saveat=0:0.5:10.0, abstol=1e-10,
    reltol=1e-12)


loss_function = function (p)
    prob = remake(prob0; p = p)
    prediction = solve(prob, RadauIIA3(); u0=convert.(eltype(p), prob0.u0), saveat = 0.0:0.5:10.0, abstol=1e-10,
        reltol = 1e-12)
    tmpdata = prediction[[1, 2], :]
    tdata = target_data[[1, 2], :]
    #Main.@infiltrate
    # Calculate squared error
    return sum(abs2, tmpdata - tdata)
end
p = [2.0, 1.2, 3.3, 1.5];

fdgrad = ForwardDiff.gradient(loss_function, p)

Zygote.gradient(loss_function,p)

[oscardssmith]

that's very odd. It's unclear to me why Radau would be any different here.

[oscardssmith]

@ChrisRackauckas or @Shreyas-Ekanathan any idea why RadauIIA5 would be different here?

[ChrisRackauckas]

Complex numbers?

[oscardssmith]

Seems very plausible. @jClugstor you might need a rule for complex matrices.

[jClugstor]

Yes, seems likely there's an issue there. Although it's strange because in the case of Complex{Dual{ ... values in A or b, the overloads aren't even being gone through, because then it isn't a DualLinearProblem.

[jClugstor]

Maybe a silly question, but when we call dolinsolve (from OrdinaryDiffEqDifferentiation) like this

  linres1 = dolinsolve(integrator, linsolve1; A = W1, b = _vec(ubuff),
                linu = _vec(dw1))

we want dw1 to end up with the same value as linres1.u, right?

[oscardssmith]

yeah this relies on some weird aliasing.

[jClugstor]

Yeah that's my issue I believe. The complex stuff was a red herring, those linear solves don't even go through the overloads since they're Complex{Dual{ ... and not Dual{Complex{ .... I'll just need to fix the aliasing up.

[jClugstor]

Ok, assuming that the MTK integration test is green, this all works with support for

• Nested Duals, so Hessian / Jacobian of gradient works
• Sparse / Float only solvers work, including KLUFactorization and UMFPACKFactorization
• anything aliased with cache.u should also be updated when the cache is solved, preserving previous behavior
• Default solvers are chosen based on the type of the lowest level number i.e. the number type that actually goes through the linear solver

But

• this does not work with problems where values are Complex{Dual. I don't think it would be difficult to implement, but for now those type of problems will just have the old behavior and have the Duals go through the solvers, so it won't change anything