requirements.txt

--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

--extra-index-url https://download.pytorch.org/whl/cu118

absl-py==2.0.0
accelerate==0.25.0
aiohttp==3.9.1
aiosignal==1.3.1
async-timeout==4.0.3
attrs==23.1.0
beartype==0.16.4
certifi==2023.11.17
charset-normalizer==3.3.2
chex==0.1.85
contourpy==1.2.0
cycler==0.12.1
datasets==2.16.0
dill==0.3.7
docstring-parser==0.15
etils==1.6.0
filelock==3.13.1
flax==0.7.5
fonttools==4.46.0
frozenlist==1.4.1
fsspec==2023.10.0
huggingface-hub==0.19.4
idna==3.6
importlib-resources==6.1.1
jax==0.4.21
jax-lorax==0.2.0
jaxlib==0.4.21+cuda11.cudnn86
Jinja2==3.1.2
kiwisolver==1.4.5
markdown-it-py==3.0.0
MarkupSafe==2.1.3
matplotlib==3.8.2
mdurl==0.1.2
ml-dtypes==0.3.1
mpmath==1.3.0
msgpack==1.0.7
multidict==6.0.4
multiprocess==0.70.15
nest-asyncio==1.5.8
networkx==3.0
numpy==1.26.2
nvidia-cublas-cu11==11.11.3.6
nvidia-cuda-cupti-cu11==11.8.87
nvidia-cuda-nvcc-cu11==11.8.89
nvidia-cuda-nvrtc-cu11==11.8.89
nvidia-cuda-runtime-cu11==11.8.89
nvidia-cudnn-cu11==8.9.6.50
nvidia-cufft-cu11==10.9.0.58
nvidia-cusolver-cu11==11.4.1.48
nvidia-cusparse-cu11==11.7.5.86
nvidia-nccl-cu11==2.19.3
opt-einsum==3.3.0
optax==0.1.7
orbax-checkpoint==0.4.7
packaging==23.2
pandas==2.1.4
Pillow==9.3.0
plum-dispatch==2.2.2
protobuf==4.25.1
psutil==5.9.7
pyarrow==14.0.2
pyarrow-hotfix==0.6
Pygments==2.17.2
pyparsing==3.1.1
python-dateutil==2.8.2
pytz==2023.3.post1
PyYAML==6.0.1
qax==0.2.0
regex==2023.10.3
requests==2.31.0
rich==13.7.0
safetensors==0.4.1
scipy==1.11.4
shtab==1.6.5
six==1.16.0
sympy==1.12
tensorstore==0.1.51
tokenizers==0.15.0
toolz==0.12.0
torch==2.1.1+cu118
torchaudio==2.1.1+cu118
torchvision==0.16.1+cu118
tqdm==4.66.1
transformers==4.36.0
triton==2.1.0
trl==0.7.7
typing_extensions==4.9.0
tyro==0.6.2
tzdata==2023.4
urllib3==2.1.0
xxhash==3.4.1
yarl==1.9.4
zipp==3.17.0