jax
jaxlib
flax
optax
numpy

[dev]
pytest
pytest-cov
hypothesis
ruff
pre-commit
mypy
mkdocs
mkdocs-material
mkdocstrings[python]
