requirements.txt

absl-py==2.1.0
chex==0.1.87
diffrax==0.6.0
einops==0.8.0
equinox==0.11.8
etils==1.10.0
filelock==3.13.1
fsspec==2024.2.0
grpcio==1.67.0
jax==0.4.35
jax-cuda12-pjrt==0.4.35
jax-cuda12-plugin==0.4.35
jaxlib==0.4.34
jaxtyping==0.2.34
Jinja2==3.1.3
lineax==0.0.7
Markdown==3.7
MarkupSafe==2.1.5
ml_dtypes==0.5.0
mpmath==1.3.0
networkx==3.2.1
numpy==2.1.2
nvidia-cublas-cu12==12.6.3.3
nvidia-cuda-cupti-cu12==12.6.80
nvidia-cuda-nvcc-cu12==12.6.77
nvidia-cuda-runtime-cu12==12.6.77
nvidia-cudnn-cu12==9.5.0.50
nvidia-cufft-cu12==11.3.0.4
nvidia-cusolver-cu12==11.7.1.2
nvidia-cusparse-cu12==12.5.4.2
nvidia-nccl-cu12==2.23.4
nvidia-nvjitlink-cu12==12.6.77
opt_einsum==3.4.0
optax==0.2.3
optimistix==0.0.9
packaging==24.1
pillow==10.2.0
protobuf==5.28.3
scipy==1.14.1
six==1.16.0
sympy==1.13.1
tensorboard==2.18.0
tensorboard-data-server==0.7.2
toolz==1.0.0
--find-links https://download.pytorch.org/whl/torch_stable.html
torch==2.3.1+cpu
torchaudio==2.3.1+cpu
torchvision==0.18.1+cpu
tqdm==4.66.5
typeguard==2.13.3
typing_extensions==4.9.0
Werkzeug==3.0.4