# If you make changes below this line, please also make the corresponding changes to `dl-cpu-requirements.txt`!

tensorflow==2.20.0; sys_platform != 'darwin' or platform_machine != 'arm64'
tensorflow-macos==2.20.0; sys_platform == 'darwin' and platform_machine == 'arm64'
tensorflow-probability==0.24.0
tf-keras==2.20.0

--extra-index-url https://download.pytorch.org/whl/cu128  # for GPU versions of torch, torchvision
--find-links https://data.pyg.org/whl/torch-2.9.0+cu128.html  # for GPU versions of torch-scatter, torch-sparse, torch-cluster, torch-spline-conv
# specifying explicit plus-notation below so pip overwrites the existing cpu verisons
torch==2.9.0+cu128
torchvision==0.24.0+cu128
torch-scatter==2.1.2+pt29cu128
torch-sparse==0.6.18+pt29cu128
torch-cluster==1.6.3+pt29cu128
torch-spline-conv==1.2.2+pt29cu128
torch-geometric==2.5.3
# Declared explicitly so GPU depsets resolve nccl from cu128 torch
# transitively rather than being pinned by the CPU-built py3.13 lock.
nvidia-nccl-cu12; platform_system == 'Linux' and platform_machine != 'aarch64'

cupy-cuda12x==13.6.0; sys_platform != 'darwin'
cudf-cu12>=24.12.0; sys_platform != 'darwin' and python_version >= '3.11'
nixl==0.4.0; sys_platform != 'darwin'

jax==0.4.33; sys_platform != 'darwin'
jaxlib==0.4.33; sys_platform != 'darwin'
jax-cuda12-plugin[cuda12]==0.4.33; sys_platform != 'darwin'