Merge branch 'main' into qqy/downgrade_ci

Author: avik-pal

.buildkite/pipeline.yml

@@ -0,0 +1,16 @@
+steps:
+  - label: "Julia 1"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-test#v1:
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 30
+    # Don't run Buildkite if the commit message includes the text [skip tests]
+    if: build.message !~ /\[skip tests\]/
+
+env:
+  GROUP: CUDA
+  JULIA_PKG_SERVER: ""

.github/workflows/CI.yml

@@ -20,7 +20,7 @@ jobs:
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
-      - uses: actions/cache@v3
+      - uses: actions/cache@v4
         env:
           cache-name: cache-artifacts
         with:
@@ -32,9 +32,11 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        with:
+          annotate: true
         env:
           GROUP: ${{ matrix.group }}
-          JULIA_NUM_THREADS: 11
+          JULIA_NUM_THREADS: "auto"
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v3
         with:

Project.toml

@@ -1,13 +1,14 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.1.0"
+version = "1.3.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -17,23 +18,30 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
-[extensions]
-SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
-
 [weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+
+[extensions]
+SimpleNonlinearSolveChainRulesCoreExt = "ChainRulesCore"
+SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
+SimpleNonlinearSolveStaticArraysExt = "StaticArrays"
 
 [compat]
 ADTypes = "0.2.6"
 ArrayInterface = "7.7"
+ChainRulesCore = "1"
 ConcreteStructs = "0.2.2"
 DiffEqBase = "6.144"
+FastClosures = "0.3"
 FiniteDiff = "2.21"
 ForwardDiff = "0.10.36"
 LinearAlgebra = "1.9"
 MaybeInplace = "0.1.1"
 PrecompileTools = "1.2"
 Reexport = "1.2"
-SciMLBase = "2.11"
+SciMLBase = "2.23"
+StaticArrays = "1"
 StaticArraysCore = "1.4.2"
-julia = "1.9"
+julia = "1.10"

README.md

@@ -5,6 +5,7 @@
 
 [![codecov](https://codecov.io/gh/SciML/SimpleNonlinearSolve.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/SciML/SimpleNonlinearSolve.jl)
 [![Build Status](https://github.com/SciML/SimpleNonlinearSolve.jl/workflows/CI/badge.svg)](https://github.com/SciML/SimpleNonlinearSolve.jl/actions?query=workflow%3ACI)
+[![Build status](https://badge.buildkite.com/c5f7db4f1b5e8a592514378b6fc807d934546cc7d5aa79d645.svg?branch=main)](https://buildkite.com/julialang/simplenonlinearsolve-dot-jl)
 
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor%27s%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
 [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle)

ext/SimpleNonlinearSolveChainRulesCoreExt.jl

@@ -0,0 +1,21 @@
+module SimpleNonlinearSolveChainRulesCoreExt
+
+using ChainRulesCore, DiffEqBase, SciMLBase, SimpleNonlinearSolve
+
+# The expectation here is that no-one is using this directly inside a GPU kernel. We can
+# eventually lift this requirement using a custom adjoint
+function ChainRulesCore.rrule(::typeof(SimpleNonlinearSolve.__internal_solve_up),
+        prob::NonlinearProblem,
+        sensealg::Union{Nothing, DiffEqBase.AbstractSensitivityAlgorithm}, u0, u0_changed,
+        p, p_changed, alg, args...; kwargs...)
+    out, ∇internal = DiffEqBase._solve_adjoint(prob, sensealg, u0, p,
+        SciMLBase.ChainRulesOriginator(), alg, args...; kwargs...)
+    function ∇__internal_solve_up(Δ)
+        ∂f, ∂prob, ∂sensealg, ∂u0, ∂p, ∂originator, ∂args... = ∇internal(Δ)
+        return (∂f, ∂prob, ∂sensealg, ∂u0, NoTangent(), ∂p, NoTangent(), ∂originator,
+            ∂args...)
+    end
+    return out, ∇__internal_solve_up
+end
+
+end

ext/SimpleNonlinearSolveStaticArraysExt.jl

@@ -0,0 +1,7 @@
+module SimpleNonlinearSolveStaticArraysExt
+
+using SimpleNonlinearSolve
+
+@inline SimpleNonlinearSolve.__is_extension_loaded(::Val{:StaticArrays}) = true
+
+end

src/SimpleNonlinearSolve.jl

@@ -3,17 +3,16 @@ module SimpleNonlinearSolve
 import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidations
 
 @recompile_invalidations begin
-    using ADTypes,
-        ArrayInterface, ConcreteStructs, DiffEqBase, Reexport, LinearAlgebra, SciMLBase
+    using ADTypes, ArrayInterface, ConcreteStructs, DiffEqBase, FastClosures, FiniteDiff,
+        ForwardDiff, Reexport, LinearAlgebra, SciMLBase
 
     import DiffEqBase: AbstractNonlinearTerminationMode,
         AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode,
         NonlinearSafeTerminationReturnCode, get_termination_mode,
-        NONLINEARSOLVE_DEFAULT_NORM, _get_tolerance
-    using FiniteDiff, ForwardDiff
+        NONLINEARSOLVE_DEFAULT_NORM
     import ForwardDiff: Dual
     import MaybeInplace: @bb, setindex_trait, CanSetindex, CannotSetindex
-    import SciMLBase: AbstractNonlinearAlgorithm, build_solution, isinplace
+    import SciMLBase: AbstractNonlinearAlgorithm, build_solution, isinplace, _unwrap_val
     import StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray, MMatrix, Size
 end
 
@@ -26,6 +25,7 @@ abstract type AbstractNewtonAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm e
 @inline __is_extension_loaded(::Val) = false
 
 include("utils.jl")
+include("linesearch.jl")
 
 ## Nonlinear Solvers
 include("nlsolve/raphson.jl")
@@ -50,13 +50,27 @@ include("ad.jl")
 ## Default algorithm
 
 # Set the default bracketing method to ITP
-function SciMLBase.solve(prob::IntervalNonlinearProblem; kwargs...)
-    return solve(prob, ITP(); kwargs...)
+SciMLBase.solve(prob::IntervalNonlinearProblem; kwargs...) = solve(prob, ITP(); kwargs...)
+function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Nothing, args...; kwargs...)
+    return solve(prob, ITP(), args...; kwargs...)
 end
 
-function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Nothing,
-        args...; kwargs...)
-    return solve(prob, ITP(), args...; kwargs...)
+# By Pass the highlevel checks for NonlinearProblem for Simple Algorithms
+function SciMLBase.solve(prob::NonlinearProblem, alg::AbstractSimpleNonlinearSolveAlgorithm,
+        args...; sensealg = nothing, u0 = nothing, p = nothing, kwargs...)
+    if sensealg === nothing && haskey(prob.kwargs, :sensealg)
+        sensealg = prob.kwargs[:sensealg]
+    end
+    new_u0 = u0 !== nothing ? u0 : prob.u0
+    new_p = p !== nothing ? p : prob.p
+    return __internal_solve_up(prob, sensealg, new_u0, u0 === nothing, new_p, p === nothing,
+        alg, args...; kwargs...)
+end
+
+function __internal_solve_up(_prob::NonlinearProblem, sensealg, u0, u0_changed, p,
+        p_changed, alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...)
+    prob = u0_changed || p_changed ? remake(_prob; u0, p) : _prob
+    return SciMLBase.__solve(prob, alg, args...; kwargs...)
 end
 
 @setup_workload begin

src/ad.jl

@@ -1,21 +1,43 @@
-function scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-    f = prob.f
+function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, <:AbstractArray},
+            iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}},
+        alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...) where {T, V, P, iip}
+    sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...)
+    dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p)
+    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats,
+        sol.original)
+end
+
+for algType in (Bisection, Brent, Alefeld, Falsi, ITP, Ridder)
+    @eval begin
+        function SciMLBase.solve(prob::IntervalNonlinearProblem{uType, iip,
+                    <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}},
+                alg::$(algType), args...; kwargs...) where {uType, T, V, P, iip}
+            sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...)
+            dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p)
+            return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode,
+                sol.stats, sol.original, left = Dual{T, V, P}(sol.left, partials),
+                right = Dual{T, V, P}(sol.right, partials))
+        end
+    end
+end
+
+function __nlsolve_ad(prob, alg, args...; kwargs...)
     p = value(prob.p)
     if prob isa IntervalNonlinearProblem
         tspan = value.(prob.tspan)
-        newprob = IntervalNonlinearProblem(f, tspan, p; prob.kwargs...)
+        newprob = IntervalNonlinearProblem(prob.f, tspan, p; prob.kwargs...)
     else
         u0 = value(prob.u0)
-        newprob = NonlinearProblem(f, u0, p; prob.kwargs...)
+        newprob = NonlinearProblem(prob.f, u0, p; prob.kwargs...)
     end
 
     sol = solve(newprob, alg, args...; kwargs...)
 
     uu = sol.u
-    f_p = scalar_nlsolve_∂f_∂p(f, uu, p)
-    f_x = scalar_nlsolve_∂f_∂u(f, uu, p)
+    f_p = __nlsolve_∂f_∂p(prob, prob.f, uu, p)
+    f_x = __nlsolve_∂f_∂u(prob, prob.f, uu, p)
 
-    z_arr = -inv(f_x) * f_p
+    z_arr = -f_x \ f_p
 
     pp = prob.p
     sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z)
@@ -30,60 +52,47 @@ function scalar_nlsolve_ad(prob, alg, args...; kwargs...)
     return sol, partials
 end
 
-function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector, <:AbstractArray},
-            false, <:Dual{T, V, P}}, alg::AbstractSimpleNonlinearSolveAlgorithm, args...;
-        kwargs...) where {T, V, P}
-    sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-    dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
-    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode)
-end
-
-function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector, <:AbstractArray},
-            false, <:AbstractArray{<:Dual{T, V, P}}},
-        alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...) where {T, V, P}
-    sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-    dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
-    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode)
-end
-
-function scalar_nlsolve_∂f_∂p(f, u, p)
-    ff = p isa Number ? ForwardDiff.derivative :
-         (u isa Number ? ForwardDiff.gradient : ForwardDiff.jacobian)
-    return ff(Base.Fix1(f, u), p)
+@inline function __nlsolve_∂f_∂p(prob, f::F, u, p) where {F}
+    if isinplace(prob)
+        __f = p -> begin
+            du = similar(u, promote_type(eltype(u), eltype(p)))
+            f(du, u, p)
+            return du
+        end
+    else
+        __f = Base.Fix1(f, u)
+    end
+    if p isa Number
+        return __reshape(ForwardDiff.derivative(__f, p), :, 1)
+    elseif u isa Number
+        return __reshape(ForwardDiff.gradient(__f, p), 1, :)
+    else
+        return ForwardDiff.jacobian(__f, p)
+    end
 end
 
-function scalar_nlsolve_∂f_∂u(f, u, p)
-    ff = u isa Number ? ForwardDiff.derivative : ForwardDiff.jacobian
-    return ff(Base.Fix2(f, p), u)
+@inline function __nlsolve_∂f_∂u(prob, f::F, u, p) where {F}
+    if isinplace(prob)
+        du = similar(u)
+        __f = (du, u) -> f(du, u, p)
+        ForwardDiff.jacobian(__f, du, u)
+    else
+        __f = Base.Fix2(f, p)
+        if u isa Number
+            return ForwardDiff.derivative(__f, u)
+        else
+            return ForwardDiff.jacobian(__f, u)
+        end
+    end
 end
 
-function scalar_nlsolve_dual_soln(u::Number, partials,
+@inline function __nlsolve_dual_soln(u::Number, partials,
         ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P}
     return Dual{T, V, P}(u, partials)
 end
 
-function scalar_nlsolve_dual_soln(u::AbstractArray, partials,
+@inline function __nlsolve_dual_soln(u::AbstractArray, partials,
         ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P}
-    return map(((uᵢ, pᵢ),) -> Dual{T, V, P}(uᵢ, pᵢ), zip(u, partials))
-end
-
-# avoid ambiguities
-for Alg in [Bisection]
-    @eval function SciMLBase.solve(prob::IntervalNonlinearProblem{uType, iip,
-                <:Dual{T, V, P}}, alg::$Alg, args...; kwargs...) where {uType, iip, T, V, P}
-        sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-        dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
-        return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode,
-            left = Dual{T, V, P}(sol.left, partials),
-            right = Dual{T, V, P}(sol.right, partials))
-    end
-    @eval function SciMLBase.solve(prob::IntervalNonlinearProblem{uType, iip,
-                <:AbstractArray{<:Dual{T, V, P}}}, alg::$Alg, args...;
-            kwargs...) where {uType, iip, T, V, P}
-        sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-        dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
-        return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode,
-            left = Dual{T, V, P}(sol.left, partials),
-            right = Dual{T, V, P}(sol.right, partials))
-    end
+    _partials = _restructure(u, partials)
+    return map(((uᵢ, pᵢ),) -> Dual{T, V, P}(uᵢ, pᵢ), zip(u, _partials))
 end

src/bracketing/bisection.jl

@@ -26,7 +26,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Bisection, args...
     left, right = prob.tspan
     fl, fr = f(left), f(right)
 
-    abstol = _get_tolerance(abstol,
+    abstol = __get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)

src/bracketing/brent.jl

@@ -13,7 +13,7 @@ function SciMLBase.solve(prob::IntervalNonlinearProblem, alg::Brent, args...;
     fl, fr = f(left), f(right)
     ϵ = eps(convert(typeof(fl), 1))
 
-    abstol = _get_tolerance(abstol,
+    abstol = __get_tolerance(nothing, abstol,
         promote_type(eltype(first(prob.tspan)), eltype(last(prob.tspan))))
 
     if iszero(fl)