chex
dm-haiku
matplotlib
numpy
optax
tensorflow
