Pin version of torch-xla to 2.5.1. (#20834)

This is needed to make it compatible with the pinned version of torch we're using.

Note that torch-xla 2.6 doesn't support GPU https://pypi.org/project/torch-xla/2.6.0/
GPU support will be coming back with 2.7.

requirements-jax-cuda.txt

@@ -3,11 +3,13 @@ tensorflow-cpu~=2.18.0
 tf2onnx
 
 # Torch cpu-only version (needed for testing).
-# TODO: Pinned to a version that is compatible with torch-xla
+# - torch is pinned to a version that is compatible with torch-xla
+# - torch-xla is pinned to a version that supports GPU (2.6 doesn't)
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch==2.5.1+cpu
 torchvision==0.20.1+cpu
 torch-xla==2.5.1;sys_platform != 'darwin'
+
 # Jax with cuda support.
 # TODO: Higher version breaks CI.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

requirements-tensorflow-cuda.txt

@@ -3,7 +3,8 @@ tensorflow[and-cuda]~=2.18.0
 tf2onnx
 
 # Torch cpu-only version (needed for testing).
-# TODO: Pinned to a version that is compatible with torch-xla
+# - torch is pinned to a version that is compatible with torch-xla
+# - torch-xla is pinned to a version that supports GPU (2.6 doesn't)
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch==2.5.1+cpu
 torchvision==0.20.1+cpu

requirements-torch-cuda.txt

@@ -3,11 +3,12 @@ tensorflow-cpu~=2.18.0
 tf2onnx
 
 # Torch with cuda support.
-# TODO: Pinned to a version that is compatible with torch-xla
+# - torch is pinned to a version that is compatible with torch-xla
+# - torch-xla is pinned to a version that supports GPU (2.6 doesn't)
 --extra-index-url https://download.pytorch.org/whl/cu121
 torch==2.5.1+cu121
 torchvision==0.20.1+cu121
-torch-xla==2.5.1
+torch-xla==2.5.1;sys_platform != 'darwin'
 
 # Jax cpu-only version (needed for testing).
 jax[cpu]

requirements.txt

@@ -5,11 +5,12 @@ tf_keras
 tf2onnx
 
 # Torch.
-# TODO: Pinned to a version that is compatible with torch-xla
+# - torch is pinned to a version that is compatible with torch-xla
+# - torch-xla is pinned to a version that supports GPU (2.6 doesn't)
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch==2.5.1
 torchvision==0.20.1
-torch-xla;sys_platform != 'darwin'
+torch-xla==2.5.1;sys_platform != 'darwin'
 
 # Jax.
 jax[cpu]