Merge branch 'master' into toLocal_wspec

Author: Hoomaaan

.bazelrc

@@ -79,18 +79,6 @@ build:native_arch_posix --host_copt=-march=native
 
 build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=1
 
-build:cuda --repo_env TF_NEED_CUDA=1
-# "sm" means we emit only cubin, which is forward compatible within a GPU generation.
-# "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
-build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
-build:cuda --@local_config_cuda//:enable_cuda
-build:cuda --define=xla_python_enable_gpu=true
-build:cuda --cxxopt=-DXLA_CUDA=1
-
-# Coverage with cuda/gcc/nvcc requires manually setting coverage flags.
-coverage:cuda --per_file_copt=third_party/.*,torch_xla/.*@--coverage
-coverage:cuda --linkopt=-lgcov
-
 build:acl --define==build_with_acl=true
 
 build:nonccl --define=no_nccl_support=true
@@ -105,7 +93,6 @@ build:tpu --define=with_tpu_support=true
 
 # Run tests serially with TPU and GPU (only 1 device is available).
 test:tpu --local_test_jobs=1
-test:cuda --local_test_jobs=1
 
 #########################################################################
 # RBE config options below.

.circleci/build.sh

@@ -50,7 +50,6 @@ source $XLA_DIR/xla_env
 export GCLOUD_SERVICE_KEY_FILE="$XLA_DIR/default_credentials.json"
 export SILO_NAME='cache-silo-ci-dev-3.8_cuda_12.1'  # cache bucket for CI
 export BUILD_CPP_TESTS='1'
-export TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_70,sm_75,compute_80,$TF_CUDA_COMPUTE_CAPABILITIES"
 build_torch_xla $XLA_DIR
 
 popd

.circleci/common.sh

@@ -112,7 +112,8 @@ function build_torch_xla() {
   # Need to uncomment the line below.
   # Currently it fails upstream XLA CI.
   # pip install plugins/cuda -v
-  pip install 'torch_xla[pallas]'
+  pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
+
   popd
 }
 

.github/scripts/run_tests.sh

@@ -55,6 +55,7 @@ function run_torch_xla_cpp_tests() {
                "test_tensor"
                # disable test_xla_backend_intf since it is flaky on upstream
                #"test_xla_backend_intf"
+               "test_xla_generator"
                "test_xla_sharding"
                "test_runtime"
                "test_status_dont_show_cpp_stacktraces"

.github/upstream/Dockerfile

@@ -15,11 +15,6 @@ ARG tpuvm=""
 # Disable CUDA for PyTorch
 ENV USE_CUDA "0"
 
-# Enable CUDA for XLA
-ENV XLA_CUDA "${cuda}"
-ENV TF_CUDA_COMPUTE_CAPABILITIES "${cuda_compute}"
-ENV TF_CUDA_PATHS "/usr/local/cuda,/usr/include,/usr"
-
 # CUDA build guidance
 ENV NVIDIA_VISIBLE_DEVICES all
 ENV NVIDIA_DRIVER_CAPABILITIES compute,utility

.github/workflows/_test.yml

@@ -140,7 +140,7 @@ jobs:
           set -x
 
           pip install expecttest unittest-xml-reporting
-          pip install 'torch_xla[pallas]'
+          pip install --pre 'torch_xla[pallas]' --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
           if [[ ! -z "$RUN_BENCHMARK_TESTS" ]]; then
             pip install -r pytorch/xla/benchmarks/requirements.txt

.github/workflows/_tpu_ci.yml

@@ -52,8 +52,8 @@ jobs:
           pip install fsspec
           pip install rich
           # jax and libtpu is needed for pallas tests.
-          pip install 'torch_xla[pallas]'
-          pip install 'torch_xla[tpu]' -f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html
+          pip install --pre 'torch_xla[pallas]' --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html'
+          pip install --pre 'torch_xla[tpu]' --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html'
           pip install --upgrade protobuf
       - name: Run Tests (${{ matrix.test_script }})
         if: inputs.has_code_changes == 'true'

BUILD

@@ -46,22 +46,6 @@ cc_binary(
     ]),
 )
 
-cc_binary(
-    name = "_XLAC_cuda_functions.so",
-    copts = [
-        "-fopenmp",
-        "-fPIC",
-    ],
-    linkopts = [
-        "-Wl,-soname,_XLAC_cuda_functions.so",
-    ],
-    linkshared = 1,
-    visibility = ["//visibility:public"],
-    deps = [
-        "//torch_xla/csrc:aten_cuda_functions",
-    ],
-)
-
 test_suite(
     name = "cpp_tests",
     # testonly = True,
@@ -72,15 +56,16 @@ test_suite(
         "//test/cpp:test_aten_xla_tensor_4",
         "//test/cpp:test_aten_xla_tensor_5",
         "//test/cpp:test_aten_xla_tensor_6",
+        "//test/cpp:test_debug_macros",
         "//test/cpp:test_ir",
         "//test/cpp:test_lazy",
         "//test/cpp:test_replication",
-        "//test/cpp:test_tensor",
-        "//test/cpp:test_xla_sharding",
         "//test/cpp:test_runtime",
         "//test/cpp:test_status_dont_show_cpp_stacktraces",
         "//test/cpp:test_status_show_cpp_stacktraces",
-        "//test/cpp:test_debug_macros",
+        "//test/cpp:test_tensor",
+        "//test/cpp:test_xla_generator",
+        "//test/cpp:test_xla_sharding",
         "//torch_xla/csrc/runtime:pjrt_computation_client_test",
         # "//torch_xla/csrc/runtime:ifrt_computation_client_test",
     ],

CONTRIBUTING.md

@@ -162,7 +162,7 @@ commands on your Linux machine directly, outside of the container.
       -f https://storage.googleapis.com/libtpu-releases/index.html
 
     # Optional: if you're using custom kernels, install pallas dependencies
-    pip install torch_xla[pallas]
+    pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
     ```
 
 1.  If you are running on a TPU VM, ensure `torch` and `torch_xla` were built and

README.md

@@ -32,10 +32,10 @@ Note: Builds are available for Python 3.8 to 3.11; please use one of the support
 # - for conda
 # conda create -n py311 python=3.11
 
-pip install torch==2.7.0 'torch_xla[tpu]==2.7.0'
+pip install torch==2.8.0 'torch_xla[tpu]==2.8.0'
 
 # Optional: if you're using custom kernels, install pallas dependencies
-pip install 'torch_xla[pallas]'
+pip install --pre torch_xla[pallas] --index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
 ```
 **As of 07/16/2025 and starting from Pytorch/XLA 2.8 release, PyTorch/XLA will 
 provide nightly and release wheels for Python 3.11 to 3.13**