GitHub

Frequently Asked Questions

How is the performance of Julia's NonlinearSolve.jl vs MATLAB's fzero?

This is addressed in a Twitter thread with the author of the improved fzero. On the test example:

import NonlinearSolve as NLS
import BenchmarkTools

const N = 100_000;
levels = 1.5 .* rand(N);
out = zeros(N);
myfun(x, lv) = x * sin(x) - lv

function f(out, levels, u0)
    for i in 1:N
        out[i] = NLS.solve(
            NLS.IntervalNonlinearProblem{false}(
                NLS.IntervalNonlinearFunction{false}(myfun), u0, levels[i]),
            NLS.Falsi()).u
    end
end

function f2(out, levels, u0)
    for i in 1:N
        out[i] = NLS.solve(
            NLS.NonlinearProblem{false}(NLS.NonlinearFunction{false}(myfun), u0, levels[i]),
            NLS.SimpleNewtonRaphson()).u
    end
end

BenchmarkTools.@btime f(out, levels, (0.0, 2.0))
BenchmarkTools.@btime f2(out, levels, 1.0)
  98.142 ms (0 allocations: 0 bytes)
  33.052 ms (0 allocations: 0 bytes)

MATLAB 2022a achieves 1.66s. Try this code yourself: we receive 0.009 seconds, or a 184x speedup.

For more information on performance of SciML, see the SciMLBenchmarks.

The solver tried to set a Dual Number in my Vector of Floats. How do I fix that?

This is a common problem that occurs if the code was not written to be generic based on the input types. For example, consider this example taken from this issue

import NonlinearSolve as NLS
import Random

function fff_incorrect(var, p)
    v_true = [1.0, 0.1, 2.0, 0.5]
    xx = [1.0, 2.0, 3.0, 4.0]
    xx[1] = var[1] - v_true[1]
    return var - v_true
end

v_true = [1.0, 0.1, 2.0, 0.5]
v_init = v_true .+ Random.randn!(similar(v_true)) * 0.1

prob_oop = NLS.NonlinearLeastSquaresProblem{false}(fff_incorrect, v_init)
try
    sol = NLS.solve(prob_oop, NLS.LevenbergMarquardt(); maxiters = 10000, abstol = 1e-8)
catch e
    @error e
end
┌ Error: MethodError(Float64, (Dual{ForwardDiff.Tag{NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.fff_incorrect), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Float64}}(0.1815917018883817,1.0,0.0,0.0,0.0),), 0x0000000000007c73)
└ @ Main faq.md:67

Essentially what happened was, NonlinearSolve checked that we can use ForwardDiff.jl to differentiate the function based on the input types. However, this function has xx = [1.0, 2.0, 3.0, 4.0] followed by a xx[1] = var[1] - v_true[1] where var might be a Dual number. This causes the error. To fix it:

  1. Specify the autodiff to be AutoFiniteDiff

    import ADTypes
sol = NLS.solve(prob_oop, NLS.LevenbergMarquardt(; autodiff = ADTypes.AutoFiniteDiff());
    maxiters = 10000, abstol = 1e-8)
    retcode: Success
u: 4-element Vector{Float64}:
 1.000000004051831
 0.09999999969774387
 1.9999999987188881
 0.499999999658765

    This worked but, Finite Differencing is not the recommended approach in any scenario.

  2. Rewrite the function to use PreallocationTools.jl or write it as

    function fff_correct(var, p)
    v_true = [1.0, 0.1, 2.0, 0.5]
    xx = eltype(var)[1.0, 2.0, 3.0, 4.0]
    xx[1] = var[1] - v_true[1]
    return xx - v_true
end

prob_oop = NLS.NonlinearLeastSquaresProblem{false}(fff_correct, v_init)
sol = NLS.solve(prob_oop, NLS.LevenbergMarquardt(); maxiters = 10000, abstol = 1e-8)
    retcode: StalledSuccess
u: 4-element Vector{Float64}:
 1.999999981739264
 0.08645347741392698
 1.9425836092527808
 0.48470661616033445

I thought NonlinearSolve.jl was type-stable and fast. But it isn't, why?

It is hard to say why your code is not fast. Take a look at the Diagnostics API to pin-point the problem. One common issue is that there is type instability.

If you are using the defaults for the autodiff and your problem is not a scalar or using static arrays, ForwardDiff will create type unstable code and lead to dynamic dispatch internally. See this simple example:

import NonlinearSolve as NLS
import InteractiveUtils
import ADTypes

f(u, p) = @. u^2 - p

prob = NLS.NonlinearProblem{false}(f, 1.0, 2.0)

InteractiveUtils.@code_warntype NLS.solve(prob, NLS.NewtonRaphson())
MethodInstance for CommonSolve.solve(::NonlinearProblem{Float64, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, ::GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, Nothing, Nothing, Nothing, Val{false}})
  from solve(prob::NonlinearProblem, args...; sensealg, u0, p, wrap, kwargs...) @ DiffEqBase ~/.julia/packages/DiffEqBase/cu9Gd/src/solve.jl:1147
Arguments
  #self#::Core.Const(CommonSolve.solve)
  prob::NonlinearProblem{Float64, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}
  args::Tuple{GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, Nothing, Nothing, Nothing, Val{false}}}
Body::SciMLBase.NonlinearSolution{Float64, 0, Float64, Float64, NonlinearProblem{Float64, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{nothing, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{nothing, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{Float64, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
1 ─ %1 = DiffEqBase.:(var"#solve#44")::Core.Const(DiffEqBase.var"#solve#44")
│   %2 = DiffEqBase.nothing::Core.Const(nothing)
│   %3 = DiffEqBase.nothing::Core.Const(nothing)
│   %4 = DiffEqBase.nothing::Core.Const(nothing)
│   %5 = DiffEqBase.Val(true)::Core.Const(Val{true}())
│   %6 = Core.NamedTuple()::Core.Const(NamedTuple())
│   %7 = Base.pairs(%6)::Core.Const(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
│   %8 = Core.tuple(%2, %3, %4, %5, %7, #self#, prob)::Tuple{Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, typeof(solve), NonlinearProblem{Float64, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}
│   %9 = Core._apply_iterate(Base.iterate, %1, %8, args)::SciMLBase.NonlinearSolution{Float64, 0, Float64, Float64, NonlinearProblem{Float64, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{nothing, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{nothing, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{Float64, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
└──      return %9

Notice that this was type-stable, since it is a scalar problem. Now what happens for static arrays

import StaticArrays

prob = NLS.NonlinearProblem{false}(f, StaticArrays.@SVector([1.0, 2.0]), 2.0)

InteractiveUtils.@code_warntype NLS.solve(prob, NLS.NewtonRaphson())
MethodInstance for CommonSolve.solve(::NonlinearProblem{StaticArraysCore.SVector{2, Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, ::GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, Nothing, Nothing, Nothing, Val{false}})
  from solve(prob::NonlinearProblem, args...; sensealg, u0, p, wrap, kwargs...) @ DiffEqBase ~/.julia/packages/DiffEqBase/cu9Gd/src/solve.jl:1147
Arguments
  #self#::Core.Const(CommonSolve.solve)
  prob::NonlinearProblem{StaticArraysCore.SVector{2, Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}
  args::Tuple{GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, Nothing, Nothing, Nothing, Val{false}}}
Body::SciMLBase.NonlinearSolution{Float64, 1, StaticArraysCore.SVector{2, Float64}, StaticArraysCore.SVector{2, Float64}, NonlinearProblem{StaticArraysCore.SVector{2, Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{nothing, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{nothing, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{StaticArraysCore.SVector{2, Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
1 ─ %1 = DiffEqBase.:(var"#solve#44")::Core.Const(DiffEqBase.var"#solve#44")
│   %2 = DiffEqBase.nothing::Core.Const(nothing)
│   %3 = DiffEqBase.nothing::Core.Const(nothing)
│   %4 = DiffEqBase.nothing::Core.Const(nothing)
│   %5 = DiffEqBase.Val(true)::Core.Const(Val{true}())
│   %6 = Core.NamedTuple()::Core.Const(NamedTuple())
│   %7 = Base.pairs(%6)::Core.Const(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
│   %8 = Core.tuple(%2, %3, %4, %5, %7, #self#, prob)::Tuple{Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, typeof(solve), NonlinearProblem{StaticArraysCore.SVector{2, Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}
│   %9 = Core._apply_iterate(Base.iterate, %1, %8, args)::SciMLBase.NonlinearSolution{Float64, 1, StaticArraysCore.SVector{2, Float64}, StaticArraysCore.SVector{2, Float64}, NonlinearProblem{StaticArraysCore.SVector{2, Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{nothing, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{nothing, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{StaticArraysCore.SVector{2, Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
└──      return %9

Again Type-Stable! Now let's try using a regular array:

prob = NLS.NonlinearProblem(f, [1.0, 2.0], 2.0)

InteractiveUtils.@code_warntype NLS.solve(prob, NLS.NewtonRaphson())
MethodInstance for CommonSolve.solve(::NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, ::GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, Nothing, Nothing, Nothing, Val{false}})
  from solve(prob::NonlinearProblem, args...; sensealg, u0, p, wrap, kwargs...) @ DiffEqBase ~/.julia/packages/DiffEqBase/cu9Gd/src/solve.jl:1147
Arguments
  #self#::Core.Const(CommonSolve.solve)
  prob::NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}
  args::Tuple{GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, Nothing, Nothing, Nothing, Val{false}}}
Body::SciMLBase.NonlinearSolution{Float64, 1, Vector{Float64}, Vector{Float64}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{nothing, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{nothing, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
1 ─ %1 = DiffEqBase.:(var"#solve#44")::Core.Const(DiffEqBase.var"#solve#44")
│   %2 = DiffEqBase.nothing::Core.Const(nothing)
│   %3 = DiffEqBase.nothing::Core.Const(nothing)
│   %4 = DiffEqBase.nothing::Core.Const(nothing)
│   %5 = DiffEqBase.Val(true)::Core.Const(Val{true}())
│   %6 = Core.NamedTuple()::Core.Const(NamedTuple())
│   %7 = Base.pairs(%6)::Core.Const(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
│   %8 = Core.tuple(%2, %3, %4, %5, %7, #self#, prob)::Tuple{Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, typeof(solve), NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}
│   %9 = Core._apply_iterate(Base.iterate, %1, %8, args)::SciMLBase.NonlinearSolution{Float64, 1, Vector{Float64}, Vector{Float64}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{nothing, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{nothing, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
└──      return %9

Ah it is still type stable. But internally since the chunksize is not statically inferred, it will be dynamic and lead to dynamic dispatch. To fix this, we directly specify the chunksize:

InteractiveUtils.@code_warntype NLS.solve(
    prob,
    NLS.NewtonRaphson(;
        autodiff = ADTypes.AutoForwardDiff(; chunksize = NLS.pickchunksize(prob.u0))
    )
)
MethodInstance for CommonSolve.solve(::NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, ::GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{2, Nothing}, Nothing, Nothing, Val{false}})
  from solve(prob::NonlinearProblem, args...; sensealg, u0, p, wrap, kwargs...) @ DiffEqBase ~/.julia/packages/DiffEqBase/cu9Gd/src/solve.jl:1147
Arguments
  #self#::Core.Const(CommonSolve.solve)
  prob::NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}
  args::Tuple{GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{2, Nothing}, Nothing, Nothing, Val{false}}}
Body::SciMLBase.NonlinearSolution{Float64, 1, Vector{Float64}, Vector{Float64}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{2, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{2, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
1 ─ %1 = DiffEqBase.:(var"#solve#44")::Core.Const(DiffEqBase.var"#solve#44")
│   %2 = DiffEqBase.nothing::Core.Const(nothing)
│   %3 = DiffEqBase.nothing::Core.Const(nothing)
│   %4 = DiffEqBase.nothing::Core.Const(nothing)
│   %5 = DiffEqBase.Val(true)::Core.Const(Val{true}())
│   %6 = Core.NamedTuple()::Core.Const(NamedTuple())
│   %7 = Base.pairs(%6)::Core.Const(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}())
│   %8 = Core.tuple(%2, %3, %4, %5, %7, #self#, prob)::Tuple{Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, typeof(solve), NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}
│   %9 = Core._apply_iterate(Base.iterate, %1, %8, args)::SciMLBase.NonlinearSolution{Float64, 1, Vector{Float64}, Vector{Float64}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}, GeneralizedFirstOrderAlgorithm{Missing, Missing, NewtonDescent{Nothing}, AutoForwardDiff{2, Nothing}, AutoFiniteDiff{Val{:forward}, Val{:forward}, Val{:hcentral}, Nothing, Nothing, Bool}, AutoForwardDiff{2, Nothing}, Val{false}}, Nothing, Nothing, SciMLBase.NLStats, NonlinearSolveBase.NonlinearSolveTrace{Val{false}, Val{false}, Nothing, NonlinearSolveBase.NonlinearSolveTracing{Val{:minimal}}, NonlinearProblem{Vector{Float64}, false, Float64, NonlinearFunction{false, SciMLBase.FullSpecialize, typeof(Main.f), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardNonlinearProblem}}}
└──      return %9

And boom! Type stable again. We always recommend picking the chunksize via NonlinearSolveBase.pickchunksize, however, if you manually specify the chunksize, it must be ≤ length of input. However, a very large chunksize can lead to excessive compilation times and slowdown.

NonlinearSolveBase.pickchunksize — Function
pickchunksize(x) = pickchunksize(length(x))
pickchunksize(x::Int)

Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length.

source