distrax==0.1.2
flax==0.7.0
transformers==4.33.1
einops==0.6.1
wandb==0.15.5
tensorflow==2.13.0
ml-collections==0.1.0
jax==0.4.11
optax==0.1.5
diffusers==0.18.2
ml-dtypes==0.2.0
