chex>=0.1.82
dm-haiku>=0.0.11
einops
flax  # For Haiku.
jaxtyping
matplotlib
tensor_annotations>=2.0
tensor_annotations_jax_stubs
# The current tensorflow probability release (v0.24.0) is broken: rely on nightly
# builds.
tf-nightly
tfp-nightly
# Optional for notebooks.
scikit-learn