torch
numpy
pixeltable
nvidia-ml-py
scikit-learn
umap-learn
torchvision
coreax

[dev]
pytest
pre-commit
ruff
mypy
types-pillow
types-tqdm

[gpu]
jax[cuda]
