jax
jaxlib
huggingface_hub
tokenizers
einops
torch
importlib_resources

[cuda12]
jax[cuda12]<0.6.0

[macos]
jax==0.5.0
jax-metal

[testing]
rwkv
transformers
ninja
tyro
optax
