jax==0.7.2
flax==0.12.0
numpy
scipy
einops
numba
optax
orbax-checkpoint

[cuda]
jax[cuda12]==0.7.2

[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
jsonargparse
