Skip to content

[WIP] feat: adopt a more JAX-like API #211

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

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

Merge branch 'main' into jax-api-changes

de77436
Select commit
Loading
Failed to load commit list.
Draft

[WIP] feat: adopt a more JAX-like API #211

Merge branch 'main' into jax-api-changes
de77436
Select commit
Loading
Failed to load commit list.
CodSpeed HQ / CodSpeed Performance Analysis succeeded Mar 26, 2026 in 0s

Performance Gate Passed

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

✅ 36 untouched benchmarks

Comparing jax-api-changes (de77436) with main (1977814)

Open in CodSpeed
View more details on CodSpeed HQ