absl-py
dm-haiku
dm-tree
git+https://github.com/deepmind/einshape
jax
numpy
optax
tqdm
typing-extensions
matplotlib
scikit-learn