[WIP] Testing out a `vmap` implementation

[avik-pal]

Example Usage

Package is private currently

using AutoBatching, SimpleNonlinearSolve

function solve_nlprob(f::F, u0::AbstractVector, p::Number, alg) where {F}
    prob = NonlinearProblem(f, u0, p)
    return solve(prob, alg).u
end

u0 = rand(3)
p = 2.0
f = (u, p) -> u .^ 2 .- p

solve_nlprob(f, u0, p, SimpleBroyden())

batched_solve_nlprob = vmap(solve_nlprob, Val((2,)))

u0_batched = rand(3, 5)

batched_solve_nlprob(f, u0_batched, p, SimpleBroyden())

Supported Algorithms

• [x] SimpleBroyden
• [x] SimpleDFSane
• [x] SimpleGaussNewton
• [x] SimpleHalley
• [x] SimpleKlement
• [x] SimpleLimitedMemoryBroyden
• [x] SimpleNewtonRaphson
• [ ] SimpleTrustRegion

Code Changes

• Conditionals are not going to be supported. To modify the code to support branches, use all or any to perform a reduction on the batched scalar

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 81.81818% with 4 lines in your changes are missing coverage. Please review.

Project coverage is 90.42%. Comparing base (6b2191f) to head (2a1eadb).

Files	Patch %	Lines
src/utils.jl	75.00%	4 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #141      +/-   ##
==========================================
- Coverage   90.58%   90.42%   -0.17%     
==========================================
  Files          23       23              
  Lines        1296     1295       -1     
==========================================
- Hits         1174     1171       -3     
- Misses        122      124       +2     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.