Tagcount ordering across precompiled packages- expected behavior?

[AdityaPandeyCN]

I'm hitting a nested ForwardDiff issue where < returns the wrong tag ordering when NonlinearSolve (outer) differentiates through OrdinaryDiffEq's Rosenbrock23 (inner ForwardDiff Jacobian). Gives MethodError: no method matching Float64(::Dual{...}) from the Dual nesting being flipped.

using OrdinaryDiffEqRosenbrock, NonlinearSolve

function ode(du, u, p, t)
    du[1] = -p[1] * u[1]
    du[2] = -u[1] - p[2] * u[2]
end

ode_f = ODEFunction{true, SciMLBase.FullSpecialize}(ode)
function solve_ode(p)
    prob = ODEProblem(ode_f, [1.0, 1.0], (0.0, 10.0), p)
    solve(prob, Rosenbrock23())
end

experiment = solve_ode([1.5, 2.0])
function resid!(du, p, nlparams)
    sol = solve_ode(p)
    du .= sol.(experiment.t, idxs=1) .- experiment[1, :]
end

res_prototype = zeros(length(experiment.t))
resid!(res_prototype, [1.0, 1.0], nothing)
nls = NonlinearFunction{true, SciMLBase.FullSpecialize}(resid!, resid_prototype=res_prototype)
prob = NonlinearLeastSquaresProblem(nls, [1.0, 1.0])
solve(prob, LevenbergMarquardt())  # errors

< returns the wrong ordering when two different packages nest ForwardDiff calls. NonlinearSolve (outer) creates its tag first, OrdinaryDiffEq's Rosenbrock23 (inner) creates its tag second, but tagcount gives the inner tag a lower number than the outer tag. This flips the Dual nesting and causes Float64(::Dual{...}) errors.

Can anyone confirm this on their behalf?

[devmotion]

Typically, this points to an error in the involved packages or the user code. For instance, sometimes problems are caused by containers that are not generic enough or a package (or a user) manually unwrapping duals using ForwardDiff.value without taking care of and considering the involved tag.

[AdityaPandeyCN]

The user function is just du[1] = -p[1] * u[1] , no containers, no ForwardDiff.value calls. The crash is at setindex! because p[1] * u[1] produces Dual{NLTag, Dual{OrdTag,...}} but the du buffer has eltype Dual{OrdTag, Dual{NLTag,...}}. Wouldn't that mean < is giving the wrong ordering? The tagcount values seem to confirm it:

using OrdinaryDiffEqRosenbrock, NonlinearSolve
using ForwardDiff: Tag, tagcount, Dual
using DiffEqBase: OrdinaryDiffEqTag
using NonlinearSolveBase: NonlinearSolveTag

NLT = Tag{NonlinearSolveTag, Float64}
OrdT = Tag{OrdinaryDiffEqTag, Dual{NLT, Float64, 2}}

@show tagcount(NLT)                  # 6
@show tagcount(OrdT)                 # 0
@show tagcount(NLT) < tagcount(OrdT) # false -shouldn't this be true?

Could this be caused by tagcount for NLT getting its value baked during precompilation while OrdT (which depends on NLT) can only be generated fresh at runtime?

[devmotion]

The simplest explanation is that the buffer has the wrong type. I don't see why the tagcount would matter here, you just can't write these duals to an incorrectly typed buffer.

[devmotion]

As expected this seems to be an issue in the SciML stack: SciML/OrdinaryDiffEq.jl#3381 Please reopen if the issue persists after SciML/OrdinaryDiffEq.jl#3381 is fixed.

[AdityaPandeyCN]

Sure, Thankyou for the assistance. Actually I was debugging that issue only (posted on slack earlier). I think something with FunctionWrappersWrappers caused this.

[ChrisRackauckas]

Typically, this points to an error in the involved packages or the user code. For instance, sometimes problems are caused by containers that are not generic enough or a package (or a user) manually unwrapping duals using ForwardDiff.value without taking care of and considering the involved tag.

No, these comments here are completely off and have nothing to do with reality. It would be good to clarify what's going on here as something like this will throw off all future discussions because there is definitely something to discuss here because of spooky action at a distance. I would expect most maintainers to have at least thought about this idea and whether this behavior is what one should be expecting from ForwardDiff rather than writing about unrelated topics in response.

The tag ordering heuristic is based on the idea that tags are created and immediately used. Type promotion in ForwardDiff is then based off of this tag order. What happens is that automatic conversion to duals that relies on this conversion then defines the conversion under the assumption that the outer tag will be a newer tag than the inner tag, and this is then used implicitly in the conversation as a sanity check. This heuristic makes sense if the derivative call creates the tag, since then the outer derivative tag is always newer than the inner derivative tag, so type promotion seems to "just work".

What breaks this case here is that NonlinearSolve.jl's precompilation runs before OrdinaryDiffEq.jl's precompilation, and thus the NonlinearSolveTag has lower precedence than the OrdinaryDiffEqTag. This means when automatic promotion is performed, you get a tag error. A "fix" for this is to then delete the precompilation in NonlinearSolve.jl https://github.com/SciML/NonlinearSolve.jl/blob/v4.19.0/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl#L45-L100, because then the first tag declaration is after the ODE one (since the ODE one will compile first), and the "SciML bug" goes away. I think we can agree that obviously has nothing to do with container types....

Now there is a way to fix this, and that's to never rely on automatic promotion and to always manually convert the inner and then the outer type. I'm not sure why the nested conversions done automatically done automatically don't work the same way, but I suspect there is some deep difficulty for making this work generically. But this means you can do stuff like SciML/OrdinaryDiffEq.jl#3414, where by manually taking a vector and converting it to the nested dual before calling the Jacobian will fix it. That seems to be a more reasonable fix than deleting some precompilations just for this behavior, so that's the direction we'll be going.

But it does call into question whether this could just be handled by making the type promotion code here more robust. I don't think most people would think that the order you precompile your packages would change which codes error with ForwardDiff and which ones don't, and this MWE then gives you a clear way to build examples which break if A precompiles before B, and examples which break if B precompiles before A. It may be valid to just say that defining tags that aren't right at the function time is just not supported behavior by ForwardDiff, it isn't defined in the API and here is an edge case you get from it. But to say it's expected while also not knowing what it is 😅 no this is a deeper one.

[devmotion]

No, these comments here are completely off and have nothing to do with reality.

Maybe not in this case. In general though, wrong container types and invalid use of ForwardDiff.value are the most likely causes for Float64(::Dual) errors. These have been coming up a lot. It's of course interesting to learn about additional causes and possible bugs, but this is the first time it comes up.

[devmotion]

Maybe #724 would fix the tagcount issue?

[ChrisRackauckas]

Yes that seems like it could be a general fix here. I don't know how that actually intersects with whatever is going on with precompilation though, but it's at least on the right track.

[ChrisRackauckas]

It looks like the core of the originating issue SciML/OrdinaryDiffEq.jl#3381, at least the hard part that isn't covered by SciML/OrdinaryDiffEq.jl#3488, is exactly #714, so the rest is a duplicate there. I think a lot of the Rosenbrock refactors that happened concurrently to the v7 hopefully make SciML/OrdinaryDiffEq.jl#3414 not required, which would be great because that is such a cludgy workaround, but I won't know until later today.