jax
jaxlib
flax
optax
jax-md
numpy

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