chex
datasets
einops
flax
jax[cuda12]
ml_collections
optax
wandb