cache construction is no longer type stable

[avik-pal]

I will generate a MWE but the DI extras is not type stable, atleast in the bruss case I tested.

[gdalle]

That is unfortunately expected, as preparation implies dynamically choosing a chunk/batch size B and then putting seeds in batches with type NTuple{B}. It is a bit similar to ForwardDiff's preparation in that sense. I have opened JuliaDiff/DifferentiationInterface.jl#534 to keep track and brainstorm

[avik-pal]

In this case I was testing with a concrete chunksize specified. Without chunksize our code was type-unstable as well.

[gdalle]

On my way to fixing it in JuliaDiff/DifferentiationInterface.jl#535

[gdalle]

Can you confirm that cache construction only used to be type-stable for sparse Jacobians? For dense Jacobians with ForwardDiff I don't see a way to get past ForwardDiff.Chunk(x), which is type-unstable

[avik-pal]

For dense it used to be type stable if AutoForwardDiff had chunksize specified, else it would be type-unstable.

[gdalle]

Okay I need to look at how you achieved type stability in the dense case. Can you point me to where that code lived?
nevermind, I figured it out

[gdalle]

Wait no I still don't have it, how do you statically set a chunk size while ensuring it always remains smaller than the vector length? Do you just let it error if the chunk is too big?

[avik-pal]

I think that case should be an error

[gdalle]

Alright I hadn't understood it that way.

[gdalle]

This issue should be fixed by DI v0.6.5, wanna check?

[avik-pal]

julia> using NonlinearSolve
Precompiling NonlinearSolve
  1 dependency successfully precompiled in 22 seconds. 157 already precompiled.

julia> prob = NonlinearProblem(NonlinearFunction{false}((u, p) -> u .* u .- p), [1.0, 1.0], 2.0)
NonlinearProblem with uType Vector{Float64}. In-place: false
u0: 2-element Vector{Float64}:
 1.0
 1.0

julia> init(prob, NewtonRaphson(; autodiff=AutoForwardDiff(; chunksize=2)))
GeneralizedFirstOrderAlgorithmCache(
    alg = NewtonRaphson(
       descent = NewtonDescent(),
       jacobian_ad = AutoForwardDiff(chunksize=2),
       forward_ad = AutoForwardDiff(chunksize=2)
    ),
    u = [1.0, 1.0],
    residual = [-1.0, -1.0],
    inf-norm(residual) = 1.0,
    nsteps = 0,
    retcode = Default
)

julia> @code_warntype init(prob, NewtonRaphson(; autodiff=AutoForwardDiff(; chunksize=2)))
MethodInstance for CommonSolve.init(::NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, ::GeneralizedFirstOrderAlgorithm{nothing, :NewtonRaphson, NoLineSearch{Bool}, Missing, NewtonDescent{Nothing, typeof(NonlinearSolve.DEFAULT_PRECS)}, AutoForwardDiff{2, Nothing}, AutoForwardDiff{2, Nothing}, Nothing})
  from init(prob::SciMLBase.AbstractDEProblem, args...; sensealg, u0, p, kwargs...) @ DiffEqBase /mnt/.julia/packages/DiffEqBase/vscDj/src/solve.jl:534
Arguments
  #self#::Core.Const(CommonSolve.init)
  prob::NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}
  args::Tuple{GeneralizedFirstOrderAlgorithm{nothing, :NewtonRaphson, NoLineSearch{Bool}, Missing, NewtonDescent{Nothing, typeof(NonlinearSolve.DEFAULT_PRECS)}, AutoForwardDiff{2, Nothing}, AutoForwardDiff{2, Nothing}, Nothing}}
Body::NonlinearSolve.GeneralizedFirstOrderAlgorithmCache{false, :LineSearch, false, Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64, Vector{Float64}, _A, GeneralizedFirstOrderAlgorithm{nothing, :NewtonRaphson, NoLineSearch{Bool}, Missing, NewtonDescent{Nothing, typeof(NonlinearSolve.DEFAULT_PRECS)}, AutoForwardDiff{2, Nothing}, AutoForwardDiff{2, Nothing}, Nothing}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, __T_jac_cache, __T_descent_cache, LineSearch.NoLineSearchCache{Float64}, Nothing, Nothing, Nothing, DiffEqBase.NonlinearTerminationModeCache{false, AbsSafeBestTerminationMode{:Inf, Int64, Base.Fix1{typeof(maximum), typeof(abs)}, Nothing, Int64, Float64}, SciMLBase.ReturnCode.T, Vector{Float64}, Float64, Float64, Float64, Float64, Vector{Float64}, Nothing, Nothing, Vector{Float64}, Int64, Vector{Float64}}, NonlinearSolve.NonlinearSolveTrace{false, false, TraceMinimal, Nothing, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}} where {_A, __T_jac_cache<:(NonlinearSolve.JacobianCache{false, _A, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Vector{Float64}, Vector{Float64}, Float64, AutoForwardDiff{2, Nothing}, DifferentiationInterfaceForwardDiffExt.ForwardDiffOneArgJacobianPrep{ForwardDiff.JacobianConfig{ForwardDiff.Tag{DifferentiationInterface.FixTail{NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Tuple{Float64}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DifferentiationInterface.FixTail{NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Tuple{Float64}}, Float64}, Float64, 2}}}}} where _A), __T_descent_cache<:(NonlinearSolve.NewtonDescentCache{false, false, Vector{Float64}, Nothing, __T_lincache, Nothing, Nothing, Nothing} where __T_lincache<:(NonlinearSolve.LinearSolverCache{__T_lincache, Nothing} where __T_lincache))}
1 ─ %1 = DiffEqBase.:(var"#init#41")::Core.Const(DiffEqBase.var"#init#41")
│   %2 = DiffEqBase.nothing::Core.Const(nothing)
│   %3 = DiffEqBase.nothing::Core.Const(nothing)
│   %4 = DiffEqBase.nothing::Core.Const(nothing)
│   %5 = Core.NamedTuple()::Core.Const(NamedTuple())
│   %6 = Base.pairs(%5)::Core.Const(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
│   %7 = Core.tuple(%2, %3, %4, %6, #self#, prob)::Tuple{Nothing, Nothing, Nothing, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, typeof(init), NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}
│   %8 = Core._apply_iterate(Base.iterate, %1, %7, args)::NonlinearSolve.GeneralizedFirstOrderAlgorithmCache{false, :LineSearch, false, Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64, Vector{Float64}, _A, GeneralizedFirstOrderAlgorithm{nothing, :NewtonRaphson, NoLineSearch{Bool}, Missing, NewtonDescent{Nothing, typeof(NonlinearSolve.DEFAULT_PRECS)}, AutoForwardDiff{2, Nothing}, AutoForwardDiff{2, Nothing}, Nothing}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, __T_jac_cache, __T_descent_cache, LineSearch.NoLineSearchCache{Float64}, Nothing, Nothing, Nothing, DiffEqBase.NonlinearTerminationModeCache{false, AbsSafeBestTerminationMode{:Inf, Int64, Base.Fix1{typeof(maximum), typeof(abs)}, Nothing, Int64, Float64}, SciMLBase.ReturnCode.T, Vector{Float64}, Float64, Float64, Float64, Float64, Vector{Float64}, Nothing, Nothing, Vector{Float64}, Int64, Vector{Float64}}, NonlinearSolve.NonlinearSolveTrace{false, false, TraceMinimal, Nothing, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}} where {_A, __T_jac_cache<:(NonlinearSolve.JacobianCache{false, _A, NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Vector{Float64}, Vector{Float64}, Float64, AutoForwardDiff{2, Nothing}, DifferentiationInterfaceForwardDiffExt.ForwardDiffOneArgJacobianPrep{ForwardDiff.JacobianConfig{ForwardDiff.Tag{DifferentiationInterface.FixTail{NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Tuple{Float64}}, Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DifferentiationInterface.FixTail{NonlinearFunction{false, SciMLBase.FullSpecialize, var"#1#2", LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing}, Tuple{Float64}}, Float64}, Float64, 2}}}}} where _A), __T_descent_cache<:(NonlinearSolve.NewtonDescentCache{false, false, Vector{Float64}, Nothing, __T_lincache, Nothing, Nothing, Nothing} where __T_lincache<:(NonlinearSolve.LinearSolverCache{__T_lincache, Nothing} where __T_lincache))}
└──      return %8

I will take a closer look later