einops>=0.8.1
jax-ai-stack
ml-collections>=1.1.0

[all]
