jax==0.11.2
flax==0.12.10
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[cuda13]==0.11.2

[docs]
sphinx>=8.2
sphinx-book-theme>=1.4
myst-parser>=2.0
sphinx-copybutton>=0.5
sphinx-design>=0.5

[test]
pytest

[wandb]
wandb==0.27.0
jsonargparse==4.41.0
