remove [tpuvm]

.github/workflows/tpu_ci.yml

@@ -39,6 +39,6 @@ jobs:
           pip install fsspec
           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
-          pip install torch_xla[tpuvm]
+          pip install torch_xla[tpu]
           cd pytorch/xla
           test/tpu/run_tests.sh

.kokoro/build.sh

@@ -17,7 +17,7 @@ install_environments
 pip install --user https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl \
   https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torchvision-nightly-cp38-cp38-linux_x86_64.whl \
   https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-nightly-cp38-cp38-linux_x86_64.whl
-pip install torch_xla[tpuvm] --user
+pip install torch_xla[tpu] --user
 
 set +e
 # -E is for preserving environment

infra/ansible/roles/build_srcs/tasks/main.yaml

@@ -19,7 +19,7 @@
 
 - name: Build PyTorch/XLA
   ansible.builtin.command:
-    cmd: python setup.py bdist_wheel
+    cmd: python setup.py bdist_wheel -p manylinux_2_28_x86_64
     chdir: "{{ (src_root, 'pytorch/xla') | path_join }}"
   environment: "{{ env_vars }}"
 

scripts/build_torch_wheels.sh

@@ -280,7 +280,7 @@ function build_and_install_torch_xla() {
   python setup.py bdist_wheel
   pip install dist/*.whl
   if [ "$TPUVM_MODE" == "1" ]; then
-    pip install torch_xla[tpuvm]
+    pip install torch_xla[tpu]
     sudo apt-get install -y google-perftools
   fi
 

setup.py

@@ -304,9 +304,6 @@ def run(self):
         # On Cloud TPU VM install with:
         # pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html
         'tpu': [f'libtpu-nightly=={_libtpu_version}'],
-        # On nightly, install libtpu with `pip install torch_xla[tpuvm]`
-        # Remove from release branches since this is not allowed by PyPI.
-        'tpuvm': [f'libtpu-nightly @ {_libtpu_storage_path}'],
         # 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=={_jax_version}', f'jax=={_jax_version}'],
     },