torch<2.15.0,>=2.14.0
torchvision<0.30.0,>=0.29.0
einops>=0.6.0
numpy>=1.21.0
pyarrow<20.0.0,>=14.0.0
datasets>=2.14.0
pytorch-lightning>=2.0.0
omegaconf>=2.3.0
wandb>=0.15.0
rich>=13.0.0
Pillow

[all]
nvsubquadratic[baselines,cuda,dali,distributed,quack]

[baselines]
timm==1.0.24

[cuda]
subquadratic-ops-torch-cu13>=0.3.0

[dali]
nvidia-dali-cuda130==1.53.0

[distributed]
megatron-core

[quack]
quack-kernels[cu13]
