absl-py
einops
etils[epath]
flax>=0.10.3
jax>=0.10.1
numpy
portpicker
