jax
numpyro
optax

[dev]
optax
pytest
pytest-check
mypy
black
isort
autoflake
sphinx
sphinx-book-theme
furo
coverage
pytest-cov
pre-commit
myst-parser
