jax
jaxlib
jaxtyping
distrax

[dev]
pytest
pre-commit
pytest-cov
