flax<0.13,>=0.12.10
jax<0.11.3,>=0.11.2
optax>=0.2.2
orbax-checkpoint>=0.12.4
numpy
pillow
opencv-python-headless
einops
grain<0.3,>=0.2.18
rich>=13
tyro
typing_extensions>=4.14
transformers>=5.16.1
jinja2
filelock
huggingface_hub

[av]
moviepy

[cuda12]
jax-cuda12-plugin[with-cuda]<=0.11.2,>=0.11.2

[cuda13]
jax-cuda13-plugin[with-cuda]<=0.11.2,>=0.11.2

[diffusers]
diffusers>=0.34.0
torch

[eval-harness]
lm_eval>=0.4.10

[gguf]
gguf==0.19.0

[guided]
outlines-core>=0.2.14

[hpo]
optuna

[inference-clients]
ollama<0.7,>=0.6.2
openai<4,>=3.8

[interop]
safetensors

[metrics]
scipy
requests

[mlflow]
mlflow

[mutation]
cosmic-ray==8.7.0

[plots]
matplotlib

[profile]
xprof!=2.23.2,<3,>=2.23.1

[quantization]
qwix==0.1.8

[serve]
starlette>=0.40
uvicorn>=0.30

[streaming]
datasets

[tensorboard]
tensorboard

[test]
pytest
ruff==0.14.3
flatbuffers==25.12.19
pytest-xdist
safetensors
torch
scipy
wandb
diffusers==0.34.0
ftfy
transformers==5.16.1
compressed-tensors==0.17.1
datasets
mlflow
tensorboard
optuna
ollama<0.7,>=0.6.2
openai<4,>=3.8
outlines-core>=0.2.14
lm_eval[hf]==0.4.13
nbclient
nbformat
ipykernel

[tfds]
tensorflow-datasets>=4.9.10

[torch]
torch

[torchax]
torch<2.15,>=2.14
torchax==0.0.14.dev20260922

[tpu]
libtpu==0.0.48.*
requests

[vision]
torchvision>=0.29

[wan]
ftfy

[wandb]
wandb
