#absl-py  # jax already depends on it
chex
dm-pix
immutabledict
jax
tensorflow
immutabledict
clu
flax
ml_collections
optax
numba
