Add requires for backward compat

Author: avik-pal

Project.toml

@@ -10,6 +10,7 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
@@ -34,10 +35,11 @@ julia = "1.6"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BenchmarkTools", "SafeTestsets", "Pkg", "Test", "StaticArrays"]
+test = ["BenchmarkTools", "SafeTestsets", "Pkg", "Test", "StaticArrays", "NNlib"]

…impleBatchedNonlinearSolveExt/broyden.jl ext/SimpleBatchedNonlinearSolveExt.jlext/SimpleBatchedNonlinearSolveExt/broyden.jl renamed to ext/SimpleBatchedNonlinearSolveExt.jl ext/SimpleBatchedNonlinearSolveExt/broyden.jl renamed to ext/SimpleBatchedNonlinearSolveExt.jl

@@ -1,3 +1,34 @@
+module SimpleBatchedNonlinearSolveExt
+
+using SimpleNonlinearSolve, SciMLBase
+isdefined(Base, :get_extension) ? (using NNlib) : (using ..NNlib)
+
+_batch_transpose(x) = reshape(x, 1, size(x)...)
+
+_batched_mul(x, y) = x * y
+
+function _batched_mul(x::AbstractArray{T, 3}, y::AbstractMatrix) where {T}
+    return dropdims(batched_mul(x, reshape(y, size(y, 1), 1, size(y, 2))); dims = 2)
+end
+
+function _batched_mul(x::AbstractMatrix, y::AbstractArray{T, 3}) where {T}
+    return batched_mul(reshape(x, size(x, 1), 1, size(x, 2)), y)
+end
+
+function _batched_mul(x::AbstractArray{T1, 3}, y::AbstractArray{T2, 3}) where {T1, T2}
+    return batched_mul(x, y)
+end
+
+function _init_J_batched(x::AbstractMatrix{T}) where {T}
+    J = ArrayInterfaceCore.zeromatrix(x[:, 1])
+    if ismutable(x)
+        J[diagind(J)] .= one(eltype(x))
+    else
+        J += I
+    end
+    return repeat(J, 1, 1, size(x, 2))
+end
+
 function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden{true}, args...;
                            abstol = nothing, reltol = nothing, maxiters = 1000, kwargs...)
     f = Base.Fix2(prob.f, prob.p)
@@ -46,3 +77,6 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden{true}, args...;
 
     return SciMLBase.build_solution(prob, alg, xₙ, fₙ; retcode = ReturnCode.MaxIters)
 end
+
+
+end

ext/SimpleBatchedNonlinearSolveExt/SimpleBatchedNonlinearSolveExt.jl

@@ -10,6 +10,18 @@ using DiffEqBase
 
 @reexport using SciMLBase
 
+if !isdefined(Base, :get_extension)
+    using Requires
+end
+
+function __init__()
+    @static if !isdefined(Base, :get_extension)
+        @require NNlib="872c559c-99b0-510c-b3b7-b6c96a88d5cd" begin
+            include("../ext/SimpleBatchedNonlinearSolveExt.jl")
+        end
+    end
+end
+
 abstract type AbstractSimpleNonlinearSolveAlgorithm <: SciMLBase.AbstractNonlinearAlgorithm end
 abstract type AbstractBracketingAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm end
 abstract type AbstractNewtonAlgorithm{CS, AD, FDT} <: AbstractSimpleNonlinearSolveAlgorithm end

ext/SimpleBatchedNonlinearSolveExt/lbroyden.jl

@@ -1,25 +1,20 @@
 """
-    Broyden()
+    Broyden(; batched = false)
 
 A low-overhead implementation of Broyden. This method is non-allocating on scalar
 and static array problems.
 """
 struct Broyden{batched} <: AbstractSimpleNonlinearSolveAlgorithm
-    Broyden(batched = false) = new{batched}()
+    Broyden(; batched = false) = new{batched}()
 end
 
 function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden{false}, args...;
                            abstol = nothing, reltol = nothing, maxiters = 1000, kwargs...)
     f = Base.Fix2(prob.f, prob.p)
     x = float(prob.u0)
 
-    # if batch && ndims(x) != 2
-    #     error("`batch` mode works only if `ndims(prob.u0) == 2`")
-    # end
-
     fₙ = f(x)
     T = eltype(x)
-    # J⁻¹ = init_J(x; batch)
     J⁻¹ = init_J(x)
 
     if SciMLBase.isinplace(prob)