jax
jaxlib
chex
typing_extensions
matplotlib
numpy
pytest
scipy
tensorflow_probability
tqdm
dm-haiku
optax
