luxai-s2==2.2.0
jax[cuda11_cudnn82]==0.4.7
chex

[dev]
toml
pre-commit
pytype
tox
tox-conda
pytest
pytest-cov
isort
autoflake
yapf
bump2version
rich
pandas
ipywidgets<8.0,>=7.0
notebook
build
twine

[torch]
torch
