From 6f5d921876ef2f389d9f07907899a19f339a2ab0 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 21 Aug 2025 09:41:26 -0400 Subject: [PATCH] Specialize functions on f parameter for trim compatibility Similar to SciML/NonlinearSolve.jl#684 and SciML/OrdinaryDiffEq.jl#2854, this PR adds type specialization for the f parameter in several functions across NonlinearSolve.jl to improve compatibility with --trim and reduce dynamic dispatch. Functions specialized: - _make_py_residual in NonlinearSolveSciPy.jl - _make_py_scalar in NonlinearSolveSciPy.jl - dogleg_method!! in trust_region.jl - construct_jacobian in jacobian_handling.jl (both general and AutoEnzyme versions) These functions either call f directly, pass f to other functions, or work with f in ways that benefit from type specialization for better compiler optimizations. --- .../src/jacobian_handling.jl | 4 ++-- lib/NonlinearSolveSciPy/src/NonlinearSolveSciPy.jl | 4 ++-- lib/SimpleNonlinearSolve/src/trust_region.jl | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/lib/NonlinearSolveHomotopyContinuation/src/jacobian_handling.jl b/lib/NonlinearSolveHomotopyContinuation/src/jacobian_handling.jl index 2c26b6be0..5165b62ba 100644 --- a/lib/NonlinearSolveHomotopyContinuation/src/jacobian_handling.jl +++ b/lib/NonlinearSolveHomotopyContinuation/src/jacobian_handling.jl @@ -111,7 +111,7 @@ and `p` the parameter object. The returned function must have the signature required by `HomotopySystemWrapper`. """ -function construct_jacobian(f, autodiff, variant, u0, p) +function construct_jacobian(f::F, autodiff, variant, u0, p) where F if variant == Scalar tmp = reinterpret(Float64, Vector{ComplexF64}(undef, 1)) else @@ -182,7 +182,7 @@ end Construct an `EnzymeJacobian` function. """ -function construct_jacobian(f, autodiff::AutoEnzyme, variant, u0, p) +function construct_jacobian(f::F, autodiff::AutoEnzyme, variant, u0, p) where F if variant == Scalar prep = DI.prepare_derivative(f, autodiff, u0, DI.Constant(p), strict = Val(false)) else diff --git a/lib/NonlinearSolveSciPy/src/NonlinearSolveSciPy.jl b/lib/NonlinearSolveSciPy/src/NonlinearSolveSciPy.jl index c328e508a..5d7165722 100644 --- a/lib/NonlinearSolveSciPy/src/NonlinearSolveSciPy.jl +++ b/lib/NonlinearSolveSciPy/src/NonlinearSolveSciPy.jl @@ -94,7 +94,7 @@ end """ Internal: wrap a Julia residual function into a Python callable """ -function _make_py_residual(f, p) +function _make_py_residual(f::F, p) where F return pyfunc(x_py -> begin x = Vector{Float64}(x_py) r = f(x, p) @@ -105,7 +105,7 @@ end """ Internal: wrap a Julia scalar function into a Python callable """ -function _make_py_scalar(f, p) +function _make_py_scalar(f::F, p) where F return pyfunc(x_py -> begin x = Float64(x_py) return f(x, p) diff --git a/lib/SimpleNonlinearSolve/src/trust_region.jl b/lib/SimpleNonlinearSolve/src/trust_region.jl index 25d34bc17..aa6041d61 100644 --- a/lib/SimpleNonlinearSolve/src/trust_region.jl +++ b/lib/SimpleNonlinearSolve/src/trust_region.jl @@ -193,7 +193,7 @@ function SciMLBase.__solve( return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters) end -function dogleg_method!!(cache, J, f, g, Δ) +function dogleg_method!!(cache, J, f::F, g, Δ) where F (; δsd, δN_δsd, δN) = cache # Compute the Newton step