numpy>=1.20
scipy>=1.7
torch
ninja

[all]
torch
jax
jaxlib

[cuda]
torch

[datasets]
requests>=2.25
tqdm>=4.60

[jax]
jax
jaxlib

[torch]
torch
