jax
jaxlib
numpy

[dev]
bump-my-version
darglint
totypes[tests]
pre-commit

[tests]
optax
parameterized
pytest
pytest-cov
pytest-subtests
