Fix adjoint for NonlinearSolution constructor

[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

Should hopefully fix SciML/NonlinearSolve.jl#581, in conjuction with #997

[ChrisRackauckas]

Tests fail.

[DhairyaLGandhi]

We should probably track .prob as well (for gradients against parameters).

Also what is the type of ȳ ? What types could we encounter here? If it's another NonlinearSolution, we should investigate what produces it and potentially try to return a tangent type

[DhairyaLGandhi]

an mwe is

@parameters g
@variables x(t) y(t) [state_priority = 10] λ(t)
eqs = [D(D(x)) ~ λ * x
       D(D(y)) ~ λ * y - g
       x^2 + y^2 ~ 1]
@mtkbuild pend = ODESystem(eqs, t)

prob = ODEProblem(pend, [x => 1, y => 0], (0.0, 1.5), [g => 1.5], guesses = [λ => 1])

sol = solve(prob, Rodas5P())

get_vars = getsym(prob, [pend.x+pend.y]);

Zygote.gradient(sol) do sol
    u = get_vars(sol)
    # u = sol[pend.x+pend.y]
    sum(reduce(vcat, u))
end

which points to this branch

SciMLBase.jl/ext/SciMLBaseZygoteExt.jl

Line 106 in 42cdf6a

VA = recursivecopy(VA)

and

SciMLBase.jl/ext/SciMLBaseZygoteExt.jl

Line 188 in 42cdf6a

VA = recursivecopy(VA)

[jClugstor]

ȳ was always a NonlinearSolution when I was testing.

[jClugstor]

Sorry, just now getting back to this.
To track .prob, would it be like this:

function ChainRulesCore.rrule(
        ::Type{<:SciMLBase.NonlinearSolution{
            T, N, uType, R, P, A, O, uType2, S, Tr}}, u, resid, prob,
        args...) where {T, N, uType, R, P, A, O, uType2, S, Tr}
    function NonlinearSolutionAdjoint(ȳ)
        (NoTangent(), ȳ.u, NoTangent(), ŷ.prob, ntuple(_ -> NoTangent(), length(args))...)
    end
    SciMLBase.NonlinearSolution{T, N, uType, R, P, A, O, uType2, S, Tr}(u, resid, prob, args...),
    NonlinearSolutionAdjoint
end

?

[DhairyaLGandhi]

Yes this looks reasonable

[DhairyaLGandhi]

As a follow up, is the type of ȳ a solution type still with the latest SciMLSensitivity

[jClugstor]

ȳ is a Tangent type with the latest SciMLSensitivity:

typeof(ȳ) = Tangent{Any}(u = 1.0, resid = ChainRulesCore.ZeroTangent(), prob = ChainRulesCore.ZeroTangent(), alg = ChainRulesCore.ZeroTangent(), retcode = ChainRulesCore.ZeroTangent(), original = ChainRulesCore.ZeroTangent(), left = ChainRulesCore.ZeroTangent(), right = ChainRulesCore.ZeroTangent(), stats = ChainRulesCore.ZeroTangent(), trace = ChainRulesCore.ZeroTangent())