jax
jaxlib
numpy
optax

[dev]
pytest
ruff
mypy
pre-commit
