wip

Author: tengyifei

.circleci/common.sh

@@ -110,11 +110,6 @@ function install_post_deps_pytorch_xla() {
   # Install dependencies after we built torch_xla. This is due to installing
   # those packages can potentially trigger `pip install torch_xla` if torch_xla
   # is not detected in the system.
-
-  # Install JAX dependency since a few tests depend on it.
-  pip install 'torch_xla[pallas]' \
-  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
 }
 
 function build_torch_xla() {

.github/workflows/_test.yml

@@ -128,10 +128,7 @@ jobs:
           set -x
 
           pip install expecttest unittest-xml-reporting
-          pip install torch_xla[pallas] \
-            -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-            -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
-          
+
           if [[ ! -z "$RUN_BENCHMARK_TESTS" ]]; then
             pip install -r pytorch/xla/benchmarks/requirements.txt
           fi

.github/workflows/_tpu_ci.yml

@@ -23,8 +23,7 @@ jobs:
           pip install --upgrade pip
           pip install fsspec
           pip install rich
-          # Jax nightly is needed for pallas tests.
-          pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
+          # libtpu is needed for pallas tests.
           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 --upgrade protobuf
       - name: Run Tests

CONTRIBUTING.md

@@ -87,10 +87,6 @@ using either VS Code or a local container:
   pip install torch_xla[tpu] \
     -f https://storage.googleapis.com/libtpu-wheels/index.html \
     -f https://storage.googleapis.com/libtpu-releases/index.html
-  # Optional: if you're using custom kernels, install pallas dependencies
-  pip install torch_xla[pallas] \
-    -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-    -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
   ```
 
 * If you are running on a TPU VM, ensure `torch` and `torch_xla` were built and 

README.md

@@ -29,11 +29,6 @@ To install PyTorch/XLA stable build in a new TPU VM:
 pip install torch~=2.6.0 'torch_xla[tpu]~=2.6.0' \
   -f https://storage.googleapis.com/libtpu-releases/index.html \
   -f https://storage.googleapis.com/libtpu-wheels/index.html
-
-# Optional: if you're using custom kernels, install pallas dependencies
-pip install 'torch_xla[pallas]' \
-  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
 ```
 
 To install PyTorch/XLA nightly build in a new TPU VM:
@@ -43,11 +38,6 @@ pip install --pre torch torchvision --index-url https://download.pytorch.org/whl
 pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev-cp310-cp310-linux_x86_64.whl' \
   -f https://storage.googleapis.com/libtpu-releases/index.html \
   -f https://storage.googleapis.com/libtpu-wheels/index.html
-
-# Optional: if you're using custom kernels, install pallas dependencies
-pip install 'torch_xla[pallas]' \
-  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
 ```
 
 ### C++11 ABI builds

docs/source/features/pallas.md

@@ -104,9 +104,5 @@ for effective memory management with KV cache.
 ## Dependencies
 
 The Pallas integration depends on JAX to function. However, not every
-JAX version is compatible with your installed PyTorch/XLA. To install
-the proper JAX:
-
-``` bash
-pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
-```
+JAX version is compatible with your installed PyTorch/XLA. PyTorch/XLA therefore
+pins specific versions of JAX in its setup script.

scripts/build_developer.sh

@@ -40,10 +40,5 @@ pip install torch_xla[tpu] \
   -f https://storage.googleapis.com/libtpu-wheels/index.html \
   -f https://storage.googleapis.com/libtpu-releases/index.html
 
-# Install Pallas dependencies
-pip install torch_xla[pallas] \
-  -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-  -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
-
 # Test that the library is installed correctly.
 python3 -c 'import torch_xla; print(torch_xla.devices()); import torchax; torchax.enable_globally()'

setup.py

@@ -369,6 +369,9 @@ def run(self):
         # importlib.metadata backport required for PJRT plugin discovery prior
         # to Python 3.10
         'importlib_metadata>=4.6;python_version<"3.10"',
+        # Some torch operations are lowered to HLO via JAX.
+        f'jaxlib=={_jaxlib_version}',
+        f'jax=={_jax_version}',
     ],
     package_data={
         'torch_xla': ['lib/*.so*',],
@@ -390,6 +393,8 @@ def run(self):
             f'libtpu=={_libtpu_version}',
             'tpu-info',
         ],
+        # As of https://github.com/pytorch/xla/pull/8895, jax is always a dependency of torch_xla.
+        # However, this no-op extras_require entrypoint is left here for backwards compatibility.
         # pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
         'pallas': [f'jaxlib=={_jaxlib_version}', f'jax=={_jax_version}'],
     },

test/tpu/xla_test_job.yaml

@@ -43,7 +43,6 @@ spec:
     - |
       pip install expecttest==0.1.6
       pip install rich
-      pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
 
       cd /src/pytorch/xla
     volumeMounts:

torchax/pyproject.toml

@@ -41,7 +41,7 @@ path = "torchax/__init__.py"
 
 [project.optional-dependencies]
 cpu = ["jax[cpu]>=0.4.30", "jax[cpu]"]
-# Add libtpu index `-f https://storage.googleapis.com/libtpu-releases/index.html`
+# Add libtpu index `-f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html`
 tpu = ["jax[cpu]>=0.4.30", "jax[tpu]"]
 cuda = ["jax[cpu]>=0.4.30", "jax[cuda12]"]
 odml = ["jax[cpu]>=0.4.30", "jax[cpu]"]