Skip to content

jax vectorisation#2636

Merged
FFroehlich merged 10 commits intodevelopfrom
jax_pmap
Jan 30, 2025
Merged

jax vectorisation#2636
FFroehlich merged 10 commits intodevelopfrom
jax_pmap

Conversation

@FFroehlich
Copy link
Member

@FFroehlich FFroehlich commented Jan 29, 2025

Looks like some recent jax update already brought massive improvements to jit time. This reduces compile time by about 4x, (~28s to ~6s on M2 MBA for Bachmann MSB2011)

@codecov
Copy link

codecov bot commented Jan 29, 2025

Codecov Report

Attention: Patch coverage is 96.77419% with 2 lines in your changes missing coverage. Please review.

Project coverage is 77.09%. Comparing base (923b8f4) to head (981cefa).
Report is 65 commits behind head on develop.

Files with missing lines Patch % Lines
python/sdist/amici/jax/petab.py 96.29% 2 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@             Coverage Diff             @@
##           develop    #2636      +/-   ##
===========================================
+ Coverage    77.06%   77.09%   +0.03%     
===========================================
  Files          332      332              
  Lines        22972    23012      +40     
  Branches      1478     1478              
===========================================
+ Hits         17704    17742      +38     
- Misses        5257     5259       +2     
  Partials        11       11              
Flag Coverage Δ
cpp 74.24% <83.87%> (-0.30%) ⬇️
cpp_python 33.52% <0.00%> (-0.40%) ⬇️
petab 39.00% <93.54%> (+0.13%) ⬆️
python 72.25% <85.48%> (-0.26%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
python/sdist/amici/jax/model.py 87.74% <100.00%> (+0.57%) ⬆️
python/sdist/amici/jax/petab.py 97.42% <96.29%> (-0.72%) ⬇️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@FFroehlich FFroehlich marked this pull request as ready for review January 29, 2025 21:48
@FFroehlich FFroehlich requested a review from a team as a code owner January 29, 2025 21:48
Copy link
Member

@dweindl dweindl left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me.

Co-authored-by: Daniel Weindl <dweindl@users.noreply.github.com>
@FFroehlich FFroehlich enabled auto-merge January 30, 2025 09:10
@sonarqubecloud
Copy link

@FFroehlich FFroehlich added this pull request to the merge queue Jan 30, 2025
Merged via the queue into develop with commit 64de63a Jan 30, 2025
32 checks passed
@dweindl dweindl deleted the jax_pmap branch June 26, 2025 16:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants