jax==0.8.1
jaxlib==0.8.1
numpy==2.3.5
optax==0.2.6

[dev]
pytest==9.0.1
ruff==0.14.5
mypy==1.18.2
pre-commit==4.4.0

[torch]
torch>=2.0.0
