[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