jax==0.7.1
flax==0.12.0
numpy==2.3.3
scipy==1.16.2
einops==0.8.1
numba==0.62.0
optax==0.2.5
orbax-checkpoint==0.11.25

[cuda]
jax[cuda12]==0.7.1

[docs]
sphinx>=7.2
pydata-sphinx-theme>=0.15
myst-parser>=2.0
sphinx-copybutton>=0.5
sphinx-design>=0.5

[test]
pytest

[wandb]
wandb==0.27.0
jsonargparse==4.41.0
