absl-py
chex
distrax
flax
jax
jaxtyping
matplotlib
numpy
optax
scipy
scikit-learn
notebook
