absl-py==1.0.0
clu==0.0.7
flax==0.5.0
jax[cpu]==0.3.13
jaxlib==0.3.10
optax @ git+https://github.com/deepmind/optax.git@7d0b8383035cda7ed1d9f2e16a81067418938a5c
tensorflow==2.9.1
tensorflow-addons==0.17.0
tensorflow_datasets==4.5.2