Put NVIDIA wheels in ci/official/requirements_updater/nvidia-requirements.txt to have central place defining wheel versions.

This change is made to have consistency in parsing NVIDIA dependencies with [JAX repository](https://github.com/jax-ml/jax/pull/30706)

`nvidia-requirements.txt` is used in Bazel hermetic python lock files and in the `tools/pip_package/setup.py` with the package requirements.

The file content is saved in the `nvidia_wheel_versions repository` and passed as an argument to `modify_setup_py.py` that populates `setup_py.tpl` script.

PiperOrigin-RevId: 802755816
This commit is contained in:
A. Unique TensorFlower 2025-09-03 16:30:48 -07:00 committed by TensorFlower Gardener
parent bf716cff71
commit 5feca55740
23 changed files with 779 additions and 165 deletions

View File

@ -80,9 +80,15 @@ tf_workspace0()
load(
"@local_xla//third_party/py:python_wheel.bzl",
"nvidia_wheel_versions_repository",
"python_wheel_version_suffix_repository",
)
nvidia_wheel_versions_repository(
name = "nvidia_wheel_versions",
versions_source = "//ci/official/requirements_updater:nvidia-requirements.txt",
)
python_wheel_version_suffix_repository(name = "tf_wheel_version_suffix")
load(

View File

@ -16,14 +16,27 @@
load("@python//:defs.bzl", "compile_pip_requirements")
load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
# TODO(ybaturina): Remove once TF is migrated to CUDA 12.9.
genrule(
name = "nvidia_constraints",
srcs = ["nvidia-requirements.txt"],
outs = ["nvidia-constraints.txt"],
cmd = """sed -E "s/>=/==/" $(location nvidia-requirements.txt) > $@;""",
)
compile_pip_requirements(
name = "requirements",
srcs = [
"nvidia-requirements.txt",
"requirements.in",
":nvidia_constraints",
],
extra_args = [
"--allow-unsafe",
"--build-isolation",
"--rebuild",
"-c $(location :nvidia_constraints)",
],
generate_hashes = True,
requirements_in = "requirements.in",
requirements_txt = REQUIREMENTS,
)

View File

@ -34,21 +34,6 @@ packaging==23.2
setuptools==78.1.1
jax==0.4.7
zstandard==0.23.0
# NVIDIA CUDA dependencies
# Note that the wheels are downloaded only when the targets in bazel command
# contain dependencies on these wheels.
nvidia-cublas-cu12 == 12.5.3.2
nvidia-cuda-cupti-cu12 == 12.5.82
nvidia-cuda-nvrtc-cu12 == 12.5.82
nvidia-cuda-runtime-cu12 == 12.5.82
nvidia-cudnn-cu12 == 9.3.0.75
nvidia-cufft-cu12 == 11.2.3.61
nvidia-curand-cu12 == 10.3.6.82
nvidia-cusolver-cu12 == 11.6.3.83
nvidia-cusparse-cu12 == 12.5.1.3
nvidia-nccl-cu12 == 2.27.7
nvidia-nvjitlink-cu12 == 12.5.82
nvidia-nvshmem-cu12>=3.2.5
# The dependencies below are needed for TF wheel testing.
tensorflow-io-gcs-filesystem==0.37.1 ; python_version <= "3.12"
libclang >= 13.0.0

View File

@ -407,68 +407,109 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
nvidia-nvshmem-cu12==3.3.9 \
--hash=sha256:2de43cbfe559e16b8e3cb777b95f1fe2ddd5c2cfd79414b09cf9cf099feba2ba \
--hash=sha256:95ba1e98189c056eb5372bd355ab714b3741a03e6de1e32f167f5240fd967c5f
# via -r ci/official/requirements_updater/requirements.in
nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -407,61 +407,99 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
@ -469,7 +507,9 @@ nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -407,61 +407,99 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
@ -469,7 +507,9 @@ nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -411,68 +411,109 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
nvidia-nvshmem-cu12==3.3.9 \
--hash=sha256:2de43cbfe559e16b8e3cb777b95f1fe2ddd5c2cfd79414b09cf9cf099feba2ba \
--hash=sha256:95ba1e98189c056eb5372bd355ab714b3741a03e6de1e32f167f5240fd967c5f
# via -r ci/official/requirements_updater/requirements.in
nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -0,0 +1,16 @@
nvidia-cublas-cu12>=12.5.3.2,<13.0
nvidia-cuda-cupti-cu12>=12.5.82,<13.0
nvidia-cuda-nvcc-cu12>=12.5.82,<13.0
nvidia-cuda-nvrtc-cu12>=12.5.82,<13.0
nvidia-cuda-runtime-cu12>=12.5.82,<13.0
# The upper bound is set for the CUDNN API compatibility.
# See
# https://docs.nvidia.com/deeplearning/cudnn/backend/latest/developer/forward-compatibility.html#cudnn-api-compatibility
nvidia-cudnn-cu12>=9.3.0.75,<10.0
nvidia-cufft-cu12>=11.2.3.61,<12.0
nvidia-curand-cu12>=10.3.6.82,<11.0
nvidia-cusolver-cu12>=11.6.3.83,<12.0
nvidia-cusparse-cu12>=12.5.1.3,<13.0
nvidia-nccl-cu12>=2.27.7,<3.0
nvidia-nvjitlink-cu12>=12.5.82,<13.0
nvidia-nvshmem-cu12>=3.2.5

View File

@ -34,21 +34,6 @@ packaging==23.2
setuptools==78.1.1
jax==0.4.7
zstandard==0.23.0
# NVIDIA CUDA dependencies
# Note that the wheels are downloaded only when the targets in bazel command
# contain dependencies on these wheels.
nvidia-cublas-cu12 == 12.5.3.2
nvidia-cuda-cupti-cu12 == 12.5.82
nvidia-cuda-nvrtc-cu12 == 12.5.82
nvidia-cuda-runtime-cu12 == 12.5.82
nvidia-cudnn-cu12 == 9.3.0.75
nvidia-cufft-cu12 == 11.2.3.61
nvidia-curand-cu12 == 10.3.6.82
nvidia-cusolver-cu12 == 11.6.3.83
nvidia-cusparse-cu12 == 12.5.1.3
nvidia-nccl-cu12 == 2.27.7
nvidia-nvjitlink-cu12 == 12.5.82
nvidia-nvshmem-cu12>=3.2.5
# The dependencies below are needed for TF wheel testing.
tensorflow-io-gcs-filesystem==0.37.1 ; python_version <= "3.12"
libclang >= 13.0.0

View File

@ -426,68 +426,109 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
nvidia-nvshmem-cu12==3.3.9 \
--hash=sha256:2de43cbfe559e16b8e3cb777b95f1fe2ddd5c2cfd79414b09cf9cf099feba2ba \
--hash=sha256:95ba1e98189c056eb5372bd355ab714b3741a03e6de1e32f167f5240fd967c5f
# via -r ci/official/requirements_updater/requirements.in
nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -426,61 +426,99 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
@ -488,7 +526,9 @@ nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -426,68 +426,109 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -426,61 +426,99 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
@ -488,7 +526,9 @@ nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -420,68 +420,109 @@ nvidia-cublas-cu12==12.5.3.2 \
--hash=sha256:7d0191251180de606023d396b94d66f66470a0ae96d1dbb906c7656ea0f71eda \
--hash=sha256:ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cudnn-cu12
# nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
--hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \
--hash=sha256:bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953 \
--hash=sha256:d32c06490c6ba35c4323730820c7d0c4c126c04ed58d2f57275adb8d54b138fe
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvcc-cu12==12.5.82 \
--hash=sha256:6eaa264da57a893ae7606dd80b169d9783444af941697822cb82c8379ffc4957 \
--hash=sha256:ab02fe922cee01235b7950f045042219fe83e15aceb4cd3c1d36db30b034dec7 \
--hash=sha256:b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-nvrtc-cu12==12.5.82 \
--hash=sha256:3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd \
--hash=sha256:5bb6a0eb01d4974bb7ca3d48bd3859472debb3c3057a5e7de2b08fbdf35eed7e \
--hash=sha256:e5db37e990056c70953b7772dd778336ef9da0a0b5bb28f9f2a61c2e42b51d78
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cuda-runtime-cu12==12.5.82 \
--hash=sha256:0fd5fbca289bceb9f0690aa9858f06187b554fdeb7e2711dfd5bb3ce58900b46 \
--hash=sha256:3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115 \
--hash=sha256:71f015dbf9df05dd71f7480132c6ebf47a6ceb2ab53d7db8e08e4b30ebb87e14
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cudnn-cu12==9.3.0.75 \
--hash=sha256:9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d \
--hash=sha256:c5cf7ff3415e446adf195a5b7dd2ba56cd00c3ee78bfdc566e51698931aa4b7f \
--hash=sha256:c819e82eed8cf564b9d37478ea4eab9e87194bb3b7f7f8098bc1f67c9b80f1b6
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cufft-cu12==11.2.3.61 \
--hash=sha256:4a8f6f0ce93c52a50ee83422a80472b5f376054a63f38532d0eab4007e7ef28b \
--hash=sha256:6d45b48a5ee7599e57131129cda2c58544d9b78b95064d3ec3e5c6b96e2b58cc \
--hash=sha256:9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-curand-cu12==10.3.6.82 \
--hash=sha256:0631ba65231260ad832ce233ddda57e7b3b7158eabf000d78e46cbb5bd5b7aae \
--hash=sha256:2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536 \
--hash=sha256:36aabeb5990297bbce3df324ea7c7c13c3aabb140c86d50ab3b23e4ec61672f1
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusolver-cu12==11.6.3.83 \
--hash=sha256:1b8b77d2fe8abe72bb722dafb708cceaeb81f1a03999477f20b33b34f46ab885 \
--hash=sha256:6224732963cba312a84c78114b9a38c4ffabb2e2a6a120923ac99ba6f895c8cf \
--hash=sha256:93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-cusparse-cu12==12.5.1.3 \
--hash=sha256:016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752 \
--hash=sha256:33520db374e2f5ebc976d6faa1852b98c398a57e6f71150fe59705928596ffd1 \
--hash=sha256:7b97fd01f0a61628af99d0efd52132fccc8c18fc5c509f13802dccf0574a19c2
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cusolver-cu12
nvidia-nccl-cu12==2.27.7 \
--hash=sha256:4617839f3bb730c3845bf9adf92dbe0e009bc53ca5022ed941f2e23fb76e6f17 \
--hash=sha256:de5ba5562f08029a19cb1cd659404b18411ed0d6c90ac5f52f30bf99ad5809aa
# via -r ci/official/requirements_updater/requirements.in
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
nvidia-nvjitlink-cu12==12.5.82 \
--hash=sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27 \
--hash=sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697 \
--hash=sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
# via
# -r ci/official/requirements_updater/requirements.in
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
# nvidia-cufft-cu12
# nvidia-cusolver-cu12
# nvidia-cusparse-cu12
nvidia-nvshmem-cu12==3.3.9 \
--hash=sha256:2de43cbfe559e16b8e3cb777b95f1fe2ddd5c2cfd79414b09cf9cf099feba2ba \
--hash=sha256:95ba1e98189c056eb5372bd355ab714b3741a03e6de1e32f167f5240fd967c5f
# via -r ci/official/requirements_updater/requirements.in
nvidia-nvshmem-cu12==3.2.5 \
--hash=sha256:2f5798d65f1a08f9878aae17cf4d3dcbfe884d1f12cf170556cd40f2be90ca96 \
--hash=sha256:e076957d5cc72e51061a04f2d46f55df477be53e8a55d0d621be08f7aefe1d00
# via
# -c ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-constraints.txt
# -r ci/official/requirements_updater/nvidia-requirements.txt
opt-einsum==3.3.0 \
--hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
--hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549

View File

@ -168,6 +168,7 @@ tf_staging/tensorflow/tools/pip_package/BUILD:
tf_staging/tensorflow/tools/pip_package/MANIFEST.in:
tf_staging/tensorflow/tools/pip_package/README:
tf_staging/tensorflow/tools/pip_package/check_load_py_test:.py
tf_staging/tensorflow/tools/pip_package/modify_setup_py:.py
tf_staging/tensorflow/tools/pip_package/pip_smoke_test:.py
tf_staging/tensorflow/tools/pip_package/setup:.py.tpl
tf_staging/tensorflow/tools/pip_package/simple_console:.py

View File

@ -1,6 +1,7 @@
# Description:
# Tools for building the TensorFlow pip package.
load("@cuda_cudart//:version.bzl", _cudart_version = "VERSION")
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
load("@local_config_syslibs//:build_defs.bzl", "if_not_system_lib")
load(
@ -13,6 +14,7 @@ load(
)
load("@local_xla//third_party/py:python_wheel.bzl", "collect_data_files", "transitive_py_deps")
load("@local_xla//xla/tsl/mkl:build_defs.bzl", "if_enable_mkl", "if_mkl", "if_mkl_ml")
load("@nvidia_wheel_versions//:versions.bzl", "NVIDIA_WHEEL_VERSIONS")
load("//tensorflow:tensorflow.bzl", "if_wheel_dependency", "if_with_tpu_support", "transitive_hdrs")
load("//tensorflow:tf_version.bzl", "TF_SEMANTIC_VERSION_SUFFIX", "TF_VERSION")
load(
@ -268,15 +270,33 @@ transitive_py_deps(
deps = COMMON_PIP_DEPS,
)
py_binary(
name = "modify_setup_py_binary",
srcs = [
"modify_setup_py.py",
],
main = "modify_setup_py.py",
deps = ["@local_xla//third_party/py:setup_py_nvidia_dependencies_util"],
)
genrule(
name = "setup_py",
srcs = ["setup.py.tpl"],
outs = ["setup.py"],
cmd = """sed -E "s/_VERSION = '0.0.0'/_VERSION = '{wheel_version}{wheel_version_suffix}'/" \
$(location setup.py.tpl) > $@;""".format(
cmd = """
$(location :modify_setup_py_binary) \\
--template_file $(location setup.py.tpl) \\
--output_file $(OUTS) \\
--nvidia_wheel_versions_data '{nvidia_wheel_versions}' \\
--tf_version "{wheel_version}{wheel_version_suffix}" \\
--cuda_version {cuda_version}
""".format(
cuda_version = _cudart_version or "12",
nvidia_wheel_versions = NVIDIA_WHEEL_VERSIONS,
wheel_version = TF_VERSION,
wheel_version_suffix = TF_SEMANTIC_VERSION_SUFFIX,
),
tools = [":modify_setup_py_binary"],
)
py_binary(
@ -435,6 +455,7 @@ py_import(
wheel_deps = if_cuda([
"@pypi_nvidia_cublas_cu12//:pkg",
"@pypi_nvidia_cuda_cupti_cu12//:pkg",
"@pypi_nvidia_cuda_nvcc_cu12//:pkg",
"@pypi_nvidia_cuda_nvrtc_cu12//:pkg",
"@pypi_nvidia_cuda_runtime_cu12//:pkg",
"@pypi_nvidia_cudnn_cu12//:pkg",

View File

@ -0,0 +1,97 @@
# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License..
# ==============================================================================
"""Modify setup.py with TensorFlow and NVIDIA wheel versions."""
import argparse
import pathlib
import third_party.py.setup_py_nvidia_dependencies_util as util
def _update_setup_with_tf_and_nvidia_wheel_versions(
template_path: pathlib.Path,
output_path: pathlib.Path,
nvidia_wheel_versions_data: str,
tf_version: str,
cuda_version: str,
):
"""Updates a setup.py template with TensorFlow and NVIDIA wheel versions.
This function reads a setup.py template file, replaces placeholder versions
for TensorFlow and various NVIDIA-related wheels based on the provided
data, and writes the result to an output file.
Args:
template_path: Path to the input setup.py.tpl template file.
output_path: Path where the modified setup.py file will be written.
nvidia_wheel_versions_data: A string containing NVIDIA wheel version data,
with each line in the format "wheel_name version_spec".
tf_version: The version string for the TensorFlow package.
cuda_version: The CUDA version string.
"""
with open(template_path) as f:
content = f.read()
content = content.replace("_VERSION = '0.0.0'", f"_VERSION = '{tf_version}'")
content = util.get_setup_py_content_with_nvidia_wheel_versions(
content, cuda_version, nvidia_wheel_versions_data
)
with open(output_path, "w") as f:
f.write(content)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--template_file",
type=pathlib.Path,
required=True,
help="Path to the setup.py.tpl template file",
)
parser.add_argument(
"--output_file",
type=pathlib.Path,
required=True,
help="Path to write the generated setup.py file",
)
parser.add_argument(
"--nvidia_wheel_versions_data",
default=None,
required=True,
help="NVIDIA wheel versions data",
)
parser.add_argument(
"--cuda_version",
type=str,
required=True,
help="The CUDA version string",
default="12",
)
parser.add_argument(
"--tf_version",
type=str,
required=True,
help="The TensorFlow package version string",
)
args = parser.parse_args()
_update_setup_with_tf_and_nvidia_wheel_versions(
args.template_file,
args.output_file,
args.nvidia_wheel_versions_data,
args.tf_version,
args.cuda_version,
)

View File

@ -54,6 +54,22 @@ from setuptools.dist import Distribution
# result for pip.
_VERSION = '0.0.0'
cuda_version = 0 # placeholder
cuda_whl_sfx = '' # placeholder
nvidia_cublas_version = '' # placeholder
nvidia_cuda_cupti_version = '' # placeholder
nvidia_cuda_nvcc_version = '' # placeholder
nvidia_cuda_runtime_version = '' # placeholder
nvidia_cudnn_version = '' # placeholder
nvidia_cufft_version = '' # placeholder
nvidia_cusolver_version = '' # placeholder
nvidia_cusparse_version = '' # placeholder
nvidia_nccl_version = '' # placeholder
nvidia_nvjitlink_version = '' # placeholder
nvidia_cuda_nvrtc_version = '' # placeholder
nvidia_curand_version = '' # placeholder
# We use the same setup.py for all tensorflow_* packages and for the nightly
# equivalents (tf_nightly_*). The package is controlled from the argument line
# when building the pip package.
@ -145,18 +161,18 @@ if collaborator_build:
EXTRA_PACKAGES = {
'and-cuda': [
# TODO(nluehr): set nvidia-* versions based on build components.
'nvidia-cublas-cu12 >= 12.5.3.2, < 13.0',
'nvidia-cuda-cupti-cu12 >= 12.5.82, < 13.0',
'nvidia-cuda-nvcc-cu12 >= 12.5.82, < 13.0',
'nvidia-cuda-nvrtc-cu12 >= 12.5.82, < 13.0',
'nvidia-cuda-runtime-cu12 >= 12.5.82, < 13.0',
'nvidia-cudnn-cu12 >= 9.3.0.75, < 10.0',
'nvidia-cufft-cu12 >= 11.2.3.61, < 12.0',
'nvidia-curand-cu12 >= 10.3.6.82, < 11.0',
'nvidia-cusolver-cu12 >= 11.6.3.83, < 12.0',
'nvidia-cusparse-cu12 >= 12.5.1.3, < 13.0',
'nvidia-nccl-cu12 >= 2.27.7, < 3.0',
'nvidia-nvjitlink-cu12 >= 12.5.82, < 13.0',
f'nvidia-cublas{cuda_whl_sfx}{nvidia_cublas_version}',
f'nvidia-cuda-cupti{cuda_whl_sfx}{nvidia_cuda_cupti_version}',
f'nvidia-cuda-nvcc{cuda_whl_sfx}{nvidia_cuda_nvcc_version}',
f'nvidia-cuda-nvrtc{cuda_whl_sfx}{nvidia_cuda_nvrtc_version}',
f'nvidia-cuda-runtime{cuda_whl_sfx}{nvidia_cuda_runtime_version}',
f'nvidia-cudnn-cu{cuda_version}{nvidia_cudnn_version}',
f'nvidia-cufft{cuda_whl_sfx}{nvidia_cufft_version}',
f'nvidia-curand{cuda_whl_sfx}{nvidia_curand_version}',
f'nvidia-cusolver{cuda_whl_sfx}{nvidia_cusolver_version}',
f'nvidia-cusparse{cuda_whl_sfx}{nvidia_cusparse_version}',
f'nvidia-nccl-cu{cuda_version}{nvidia_nccl_version}',
f'nvidia-nvjitlink{cuda_whl_sfx}{nvidia_nvjitlink_version}',
],
'gcs-filesystem': [
('tensorflow-io-gcs-filesystem>=0.23.1; '

View File

@ -65,3 +65,9 @@ py_binary(
main = "unpack_wheel_and_unzip_archive_files.py",
visibility = ["//visibility:public"],
)
py_library(
name = "setup_py_nvidia_dependencies_util",
srcs = ["setup_py_nvidia_dependencies_util.py"],
visibility = ["//visibility:public"],
)

View File

@ -25,6 +25,7 @@ tensorflow/third_party/py/python_repo.bzl:
tensorflow/third_party/py/python_wheel.bzl:
tensorflow/third_party/py/rules_pywrap/def_file_filter_tool.py:
tensorflow/third_party/py/rules_pywrap/wrapped_py_init.cc:
tensorflow/third_party/py/setup_py_nvidia_dependencies_util.py:
tensorflow/third_party/py/unpack_wheel_and_unzip_archive_files.py:
tensorflow/tools/def_file_filter/BUILD.tpl:
tensorflow/tools/def_file_filter/BUILD:

View File

@ -65,3 +65,9 @@ py_binary(
main = "unpack_wheel_and_unzip_archive_files.py",
visibility = ["//visibility:public"],
)
py_library(
name = "setup_py_nvidia_dependencies_util",
srcs = ["setup_py_nvidia_dependencies_util.py"],
visibility = ["//visibility:public"],
)

View File

@ -257,3 +257,23 @@ It recursively traverses `deps` attribute of the target and collects paths to
files that are in `data` attribute. Then it filters all files that do not match
the provided extensions.
""" # buildifier: disable=no-effect
def _nvidia_wheel_versions_repository_impl(repository_ctx):
"""Repository rule for storing NVIDIA wheel versions."""
versions_source = repository_ctx.attr.versions_source
versions_file_content = repository_ctx.read(
repository_ctx.path(versions_source),
)
repository_ctx.file(
"versions.bzl",
"NVIDIA_WHEEL_VERSIONS = '''%s'''" % versions_file_content,
)
repository_ctx.file("BUILD", "")
nvidia_wheel_versions_repository = repository_rule(
implementation = _nvidia_wheel_versions_repository_impl,
attrs = {
"versions_source": attr.label(mandatory = True, allow_single_file = True),
},
)

View File

@ -0,0 +1,76 @@
# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility function for updating setup.py with NVIDIA wheel versions.
The content of the setup.py file is updated with the NVIDIA wheel versions
provided in the nvidia_wheel_versions_data string.
The setup.py file is expected to have the following lines:
```
# Mandatory placeholders
cuda_version = 0 # placeholder
cuda_whl_sfx = '' # placeholder
# Optional placeholders (add only those that are needed)
nvidia_cublas_version = '' # placeholder
EXTRA_PACKAGES = {
'and-cuda': [
f'nvidia-cublas{cuda_whl_sfx}{nvidia_cublas_version}',
# add more wheels here
],
}
```
"""
import re
# Regex to capture wheel name and its version constraint
# Example: "nvidia-cublas-cu12>=12.1.3.1 ; sys_platform == 'linux'"
NVIDIA_WHEEL_VERSIONS_PATTERN = re.compile(r"^([a-z0-9_-]+)(\W*[0-9\.]*.*)$")
def get_setup_py_content_with_nvidia_wheel_versions(
setup_py_content: str, cuda_version: str, nvidia_wheel_versions_data: str
) -> str:
nvidia_wheel_versions = {"12": {}, "13": {}}
for line in nvidia_wheel_versions_data.splitlines():
match = NVIDIA_WHEEL_VERSIONS_PATTERN.match(line)
if match:
wheel_name = match.group(1).replace("-", "_")
for suffix, version in {"_cu12": "12", "_cu13": "13", "": "13"}.items():
if not wheel_name.endswith(suffix):
continue
wheel_name = wheel_name.replace(suffix, "") + "_version"
nvidia_wheel_versions[version][wheel_name] = match.group(2).strip()
break
setup_py_content = setup_py_content.replace(
"cuda_version = 0 # placeholder", f"cuda_version = {cuda_version}"
)
setup_py_content = setup_py_content.replace(
"cuda_whl_sfx = '' # placeholder",
"cuda_whl_sfx = '-cu12'" if cuda_version == "12" else "cuda_whl_sfx = ''",
)
for version_name, version_value in nvidia_wheel_versions[
str(cuda_version)
].items():
setup_py_content = setup_py_content.replace(
f"{version_name} = '' # placeholder",
f"{version_name} = '{version_value}'",
)
return setup_py_content