chex
jax
jaxline
optax
dm-haiku
dm-tree
typing_extensions
matplotlib
mediapy
opencv-python
einshape
ipympl
tqdm
