Skip to content
This repository has been archived by the owner on Oct 31, 2024. It is now read-only.

[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 •
edited
Loading

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

@codecov Codecov
Copy link

codecov bot commented May 4, 2024 •
edited
Loading

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
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@avik-pal