Skip to content
New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

[WIP] Testing out a vmap implementation #141

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

[WIP] Testing out a vmap implementation #141

wants to merge 4 commits into from

Conversation

avik-pal
Copy link
Member

@avik-pal avik-pal commented May 4, 2024

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

  • SimpleBroyden
  • SimpleDFSane
  • SimpleGaussNewton
  • SimpleHalley
  • SimpleKlement
  • SimpleLimitedMemoryBroyden
  • 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

Copy link

codecov bot commented May 4, 2024

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 ⚠️
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     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@avik-pal avik-pal force-pushed the ap/batching branch 3 times, most recently from caf295c to 015b9e0 Compare May 4, 2024 22:33
# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant