diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml
index 59fd443a3668..dda9626af13a 100644
--- a/.github/workflows/wheel_tests_continuous.yml
+++ b/.github/workflows/wheel_tests_continuous.yml
@@ -35,8 +35,8 @@ permissions:
   actions: read
 
 on:
-  schedule:
-    - cron: "0 */3 * * *" # Run once every 3 hours
+  # schedule:
+  #  - cron: "0 */3 * * *" # Run once every 3 hours
   workflow_dispatch: # allows triggering the workflow run manually
 
 concurrency:
diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml
index 992f42caa084..667ed49bf30e 100644
--- a/.github/workflows/wheel_tests_nightly_release.yml
+++ b/.github/workflows/wheel_tests_nightly_release.yml
@@ -47,132 +47,6 @@ env:
   UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
 
 jobs:
-  run-pytest-cpu:
-    uses: ./.github/workflows/pytest_cpu.yml
-    strategy:
-        fail-fast: false # don't cancel all jobs on failure
-        matrix:
-          # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
-          # that build the wheels.
-          runner: ["linux-x86-n4-64", "linux-arm64-t2a-48", "windows-x86-n2-64"]
-          python: ["3.11", "3.12", "3.13", "3.13-nogil", "3.14", "3.14-nogil"]
-          enable-x64: [0]
-          exclude:
-            - runner: "windows-x86-n2-64"
-              python: "3.13-nogil"
-            - runner: "windows-x86-n2-64"
-              python: "3.14-nogil"
-    name: "Pytest CPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
-    with:
-      runner: ${{ matrix.runner }}
-      python: ${{ matrix.python }}
-      enable-x64:  ${{ matrix.enable-x64 }}
-      skip-download-jaxlib-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
-      gcs_download_uri: ${{inputs.gcs_download_uri}}
-      halt-for-connection: ${{inputs.halt-for-connection}}
-      # TODO(b/456203132): Increase to more than 16 processes if this works
-      #  (56 prior to this).
-      max-processes: ${{ contains(matrix.runner, 'windows-x86') && '16' || '' }}
-
-  run-bazel-test-cpu:
-    uses: ./.github/workflows/bazel_cpu.yml
-    strategy:
-        fail-fast: false # don't cancel all jobs on failure
-        matrix:
-          # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
-          # that build the wheels.
-          runner: ["linux-x86-n4-64", "linux-arm64-t2a-48"]
-          python: ["3.11", "3.12", "3.13", "3.13-nogil", "3.14", "3.14-nogil"]
-          enable-x64: [0]
-    name: "Bazel CPU tests with ${{ format('{0}', 'build_jaxlib=false') }}"
-    with:
-      runner: ${{ matrix.runner }}
-      python: ${{ matrix.python }}
-      enable-x64:  ${{ matrix.enable-x64 }}
-      gcs_download_uri: ${{inputs.gcs_download_uri}}
-      halt-for-connection: ${{inputs.halt-for-connection}}
-      build_jaxlib: ${{ 'false' }}
-      build_jax: ${{ 'false' }}
-      clone_main_xla: 0
-
-  run-pytest-cuda:
-    uses: ./.github/workflows/pytest_cuda.yml
-    strategy:
-        fail-fast: false # don't cancel all jobs on failure
-        matrix:
-          # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
-          # that build the wheels.
-          runner: ["linux-x86-g2-48-l4-4gpu",  "linux-x86-a3-8g-h100-8gpu", "linux-x86-a4-224-b200-1gpu"]
-          python: ["3.11", "3.12", "3.13", "3.13-nogil", "3.14", "3.14-nogil"]
-          cuda: [
-            {cuda-version: "12.1", use-nvidia-pip-wheels: false},
-            {cuda-version: "12.9", use-nvidia-pip-wheels: true},
-            {cuda-version: "13", use-nvidia-pip-wheels: true}
-          ]
-          enable-x64: [0]
-          exclude:
-            # H100 runs only CUDA 12.9 and min and max Python versions.
-            - runner: "linux-x86-a3-8g-h100-8gpu"
-              cuda:
-                cuda-version: "12.1"
-            - runner: "linux-x86-a3-8g-h100-8gpu"
-              python: "3.12"
-            - runner: "linux-x86-a3-8g-h100-8gpu"
-              python: "3.13"
-            - runner: "linux-x86-a3-8g-h100-8gpu"
-              python: "3.13-nogil"
-            - runner: "linux-x86-a3-8g-h100-8gpu"
-              python: "3.14-nogil"
-            # B200 runs only CUDA 12.9 and min and max Python versions.
-            - runner: "linux-x86-a4-224-b200-1gpu"
-              cuda:
-                cuda-version: "12.1"
-            - runner: "linux-x86-a4-224-b200-1gpu"
-              python: "3.12"
-            - runner: "linux-x86-a4-224-b200-1gpu"
-              python: "3.13"
-            - runner: "linux-x86-a4-224-b200-1gpu"
-              python: "3.13-nogil"
-            - runner: "linux-x86-a4-224-b200-1gpu"
-              python: "3.14-nogil"
-    name: "Pytest CUDA (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }}, CUDA Pip packages = ${{ matrix.cuda.use-nvidia-pip-wheels }})"
-    with:
-      runner: ${{ matrix.runner }}
-      python: ${{ matrix.python }}
-      cuda-version:  ${{ matrix.cuda.cuda-version }}
-      use-nvidia-pip-wheels: ${{ matrix.cuda.use-nvidia-pip-wheels }}
-      enable-x64:  ${{ matrix.enable-x64 }}
-      skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
-      gcs_download_uri: ${{inputs.gcs_download_uri}}
-      halt-for-connection: ${{inputs.halt-for-connection}}
-
-  run-bazel-test-cuda:
-    uses: ./.github/workflows/bazel_cuda.yml
-    strategy:
-        fail-fast: false # don't cancel all jobs on failure
-        matrix:
-          # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
-          # that build the wheels.
-          runner: ["linux-x86-g2-48-l4-4gpu"]
-          python: ["3.11", "3.12", "3.13", "3.13-nogil", "3.14", "3.14-nogil"]
-          cuda-version: [12, 13]
-          enable-x64: [0]
-    name: "Bazel CUDA Non-RBE with ${{ format('{0}', 'build_jaxlib=false') }}"
-    with:
-      runner: ${{ matrix.runner }}
-      python: ${{ matrix.python }}
-      cuda-version: ${{ matrix.cuda-version }}
-      enable-x64:  ${{ matrix.enable-x64 }}
-      gcs_download_uri: ${{inputs.gcs_download_uri}}
-      halt-for-connection: ${{inputs.halt-for-connection}}
-      build_jaxlib: ${{ 'false' }}
-      build_jax: ${{ 'false' }}
-      jaxlib-version: "head"
-      write_to_bazel_remote_cache: 1
-      run_tests_with_rbe: ${{ 'false' }}
-      build_artifacts_with_rbe: ${{ 'false' }}
-      clone_main_xla: 0
-
   run-pytest-rocm:
     permissions:
       id-token: write
@@ -198,9 +72,6 @@ jobs:
       gcs_download_uri: ${{inputs.gcs_download_uri}}
 
   run-bazel-test-rocm:
-    permissions:
-      id-token: write
-      contents: read
     uses: ./.github/workflows/bazel_rocm.yml
     strategy:
         fail-fast: false # don't cancel all jobs on failure
@@ -208,7 +79,7 @@ jobs:
           # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
           # that build the wheels.
           runner: ["linux-x86-64-1gpu-amd"]
-          python: ["3.11", "3.12", "3.13", "3.14"]
+          python: ["3.11", "3.12", "3.13", "3.13-nogil", "3.14", "3.14-nogil"]
           rocm-version: [7]
           enable-x64: [0]
     name: "Bazel ROCm Non-RBE with ${{ format('{0}', 'build_jaxlib=false') }}"
@@ -219,149 +90,11 @@ jobs:
       enable-x64:  ${{ matrix.enable-x64 }}
       gcs_download_uri: ${{inputs.gcs_download_uri}}
       halt-for-connection: ${{inputs.halt-for-connection}}
-      build_jaxlib: ${{ 'false' }}
-      build_jax: ${{ 'false' }}
+      build_jaxlib: "false"
+      build_jax: "false"
       jaxlib-version: "head"
       write_to_bazel_remote_cache: 1
-      run_multiaccelerator_tests: ${{ 'true' }}
-      clone_main_xla: 0
-
-  run-pytest-tpu:
-    uses: ./.github/workflows/pytest_tpu.yml
-    strategy:
-        fail-fast: false # don't cancel all jobs on failure
-        matrix:
-          python: ["3.11", "3.12", "3.13", "3.13-nogil", "3.14", "3.14-nogil"]
-          tpu-specs: [
-            # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
-            {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
-            {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"},
-            {type: "v7x-8", cores: "4", runner: "linux-x86-tpu7x-224-4tpu"}
-          ]
-          libtpu-version-type: ["pinned", "nightly"]
-          exclude:
-            # Exclude nightly for releases
-            - libtpu-version-type: ${{ startsWith(github.ref_name, 'release/') && 'nightly' }}
-            # Exclude pinned for nightly releases
-            - libtpu-version-type: ${{ !startsWith(github.ref_name, 'release/') && 'pinned' }}
-            # Run Python versions in between min and max for v6e-8
-            - tpu-specs:
-                type: "v6e-8"
-              python: "3.11"
-            - tpu-specs:
-                type: "v6e-8"
-              python: "3.14"
-            - tpu-specs:
-                type: "v6e-8"
-              python: "3.14-nogil"
-            # Run max Python versions for v5e-8
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.11"
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.12"
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.13"
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.13-nogil"
-            # Run min and max Python versions for v7x-8
-            - tpu-specs:
-                type: "v7x-8"
-              python: "3.12"
-            - tpu-specs:
-                type: "v7x-8"
-              python: "3.13"
-            - tpu-specs:
-                type: "v7x-8"
-              python: "3.13-nogil"
-
-    name: "Pytest TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
-    with:
-      runner: ${{ matrix.tpu-specs.runner }}
-      cores: ${{ matrix.tpu-specs.cores }}
-      tpu-type: ${{ matrix.tpu-specs.type }}
-      python: ${{ matrix.python }}
-      run-full-tpu-test-suite: "1"
-      libtpu-version-type: ${{ matrix.libtpu-version-type }}
-      skip-download-jaxlib-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
-      gcs_download_uri: ${{inputs.gcs_download_uri}}
-      halt-for-connection: ${{inputs.halt-for-connection}}
-
-  run-bazel-test-tpu:
-    uses: ./.github/workflows/bazel_test_tpu.yml
-    strategy:
-        fail-fast: false # don't cancel all jobs on failure
-        matrix:
-          python: ["3.11", "3.12", "3.13", "3.13-nogil", "3.14", "3.14-nogil"]
-          tpu-specs: [
-            # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
-            {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
-            {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
-            {type: "v7x-8", cores: "4", runner: "linux-x86-tpu7x-224-4tpu"},
-          ]
-          libtpu-version-type: ["pinned", "nightly"]
-          exclude:
-            # Exclude nightly for releases
-            - libtpu-version-type: ${{ startsWith(github.ref_name, 'release/') && 'nightly' }}
-            # Exclude pinned for nightly releases
-            - libtpu-version-type: ${{ !startsWith(github.ref_name, 'release/') && 'pinned' }}
-            # Run a single Python version for v4-8
-            - tpu-specs:
-                type: "v4-8"
-              python: "3.11"
-            - tpu-specs:
-                type: "v4-8"
-              python: "3.12"
-            - tpu-specs:
-                type: "v4-8"
-              python: "3.13"
-            - tpu-specs:
-                type: "v4-8"
-              python: "3.13-nogil"
-            - tpu-specs:
-                type: "v4-8"
-              python: "3.14-nogil"
-            # Run max Python versions for v5e-8
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.11"
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.12"
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.13"
-            - tpu-specs:
-                type: "v5e-8"
-              python: "3.13-nogil"
-            # Run min and max Python versions for v7x-8
-            - tpu-specs:
-                type: "v7x-8"
-              python: "3.12"
-            - tpu-specs:
-                type: "v7x-8"
-              python: "3.13"
-            - tpu-specs:
-                type: "v7x-8"
-              python: "3.13-nogil"
-
-    name: "Bazel tests TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
-    with:
-      runner: ${{ matrix.tpu-specs.runner }}
-      cores: ${{ matrix.tpu-specs.cores }}
-      tpu-type: ${{ matrix.tpu-specs.type }}
-      python: ${{ matrix.python }}
-      run-full-tpu-test-suite: "1"
-      libtpu-version-type: ${{ matrix.libtpu-version-type }}
-      skip-download-jaxlib-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
-      gcs_download_uri: ${{inputs.gcs_download_uri}}
-      halt-for-connection: ${{inputs.halt-for-connection}}
-      build_jaxlib: ${{ 'false' }}
-      build_jax: ${{ 'false' }}
-      jaxlib-version: "head"
+      run_multiaccelerator_tests: "true"
       clone_main_xla: 0
 
   verify-release-wheels-install:
@@ -387,85 +120,75 @@ jobs:
     steps:
       - name: Download release wheels from GCS
         run: |
-          mkdir -p "$(pwd)/dist"
-          final_gcs_download_uri="${INPUTS_GCS_DOWNLOAD_URI}"
+          mkdir -p $(pwd)/dist
+          final_gcs_download_uri=${INPUTS_GCS_DOWNLOAD_URI}
 
           # Get the major and minor version of Python.
           # E.g if python=3.11, then python_major_minor=311
           # E.g if python=3.13-nogil, then python_major_minor=313t
-          python_major_minor='${{ matrix.python }}'
-          python_major_minor="$(echo "${python_major_minor//-nogil/t}" | tr -d '.')"
+          python_major_minor=${{ matrix.python }}
+          python_major_minor=$(echo "${python_major_minor//-nogil/t}" | tr -d '.')
           python_major_minor="cp${python_major_minor%t}-cp${python_major_minor}-"
 
-          gcloud storage cp -r "${final_gcs_download_uri}"/jax*py3*none*any.whl "$(pwd)/dist/"
+          gcloud storage cp -r "${final_gcs_download_uri}"/jax*py3*none*any.whl $(pwd)/dist/
 
-          jax_wheel="$(ls dist/jax*py3*none*any.whl 2>/dev/null)"
-          {
-            echo "JAX_WHEEL=$jax_wheel"
-          } >> "$GITHUB_ENV"
+          jax_wheel=$(ls dist/jax*py3*none*any.whl 2>/dev/null)
+          echo "JAX_WHEEL=$jax_wheel" >> $GITHUB_ENV
 
           if [[ "${INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS}" != "1" ]]; then
-            gcloud storage cp -r "${final_gcs_download_uri}/jaxlib*${python_major_minor}*linux*x86_64*.whl" "$(pwd)/dist/"
-            gcloud storage cp -r "${final_gcs_download_uri}/jax*cuda${{ matrix.cuda-version }}*plugin*${python_major_minor}*linux*x86_64*.whl" "$(pwd)/dist/"
-            gcloud storage cp -r "${final_gcs_download_uri}/jax*cuda${{ matrix.cuda-version }}*pjrt*linux*x86_64*.whl" "$(pwd)/dist/"
+            gcloud storage cp -r "${final_gcs_download_uri}/jaxlib*${python_major_minor}*linux*x86_64*.whl" $(pwd)/dist/
+            gcloud storage cp -r "${final_gcs_download_uri}/jax*cuda${{ matrix.cuda-version }}*plugin*${python_major_minor}*linux*x86_64*.whl" $(pwd)/dist/
+            gcloud storage cp -r "${final_gcs_download_uri}/jax*cuda${{ matrix.cuda-version }}*pjrt*linux*x86_64*.whl" $(pwd)/dist/
 
-            jaxlib_wheel="$(ls dist/jaxlib*"${python_major_minor}"*linux*x86_64*.whl 2>/dev/null)"
-            jax_cuda_plugin_wheel="$(ls dist/jax*cuda${{ matrix.cuda-version }}*plugin*"${python_major_minor}"*linux*x86_64*.whl 2>/dev/null)"
-            jax_cuda_pjrt_wheel="$(ls dist/jax*cuda${{ matrix.cuda-version }}*pjrt*linux*x86_64*.whl 2>/dev/null)"
+            jaxlib_wheel=$(ls dist/jaxlib*${python_major_minor}*linux*x86_64*.whl 2>/dev/null)
+            jax_cuda_plugin_wheel=$(ls dist/jax*cuda${{ matrix.cuda-version }}*plugin*${python_major_minor}*linux*x86_64*.whl 2>/dev/null)
+            jax_cuda_pjrt_wheel=$(ls dist/jax*cuda${{ matrix.cuda-version }}*pjrt*linux*x86_64*.whl 2>/dev/null)
 
-            {
-              echo "JAXLIB_WHEEL=$jaxlib_wheel"
-              echo "JAX_CUDA_PLUGIN_WHEEL=$jax_cuda_plugin_wheel"
-              echo "JAX_CUDA_PJRT_WHEEL=$jax_cuda_pjrt_wheel"
-            } >> "$GITHUB_ENV"
+            echo "JAXLIB_WHEEL=$jaxlib_wheel" >> $GITHUB_ENV
+            echo "JAX_CUDA_PLUGIN_WHEEL=$jax_cuda_plugin_wheel" >> $GITHUB_ENV
+            echo "JAX_CUDA_PJRT_WHEEL=$jax_cuda_pjrt_wheel" >> $GITHUB_ENV
           fi
         env:
           INPUTS_GCS_DOWNLOAD_URI: ${{ inputs.gcs_download_uri }}
           INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
       - name: Verify JAX CPU packages can be installed
         run: |
-          # shellcheck disable=SC1090
-          "$PYTHON" -m uv venv ~/test_cpu && source ~/test_cpu/bin/activate
+          $PYTHON -m uv venv ~/test_cpu && source ~/test_cpu/bin/activate
           if [[ "${INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS}" == "1" ]]; then
-            uv pip install "$JAX_WHEEL"
+            uv pip install $JAX_WHEEL
           else
-            uv pip install "$JAX_WHEEL" "$JAXLIB_WHEEL"
+            uv pip install $JAX_WHEEL $JAXLIB_WHEEL
           fi
         env:
           INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
       - name: Verify JAX TPU packages can be installed
         run: |
-          # shellcheck disable=SC1090
-          "$PYTHON" -m uv venv ~/test_tpu && source ~/test_tpu/bin/activate
+          $PYTHON -m uv venv ~/test_tpu && source ~/test_tpu/bin/activate
 
           if [[ "${INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS}" == "1" ]]; then
-            uv pip install "${JAX_WHEEL}[tpu]"
+            uv pip install $JAX_WHEEL[tpu]
           else
-            uv pip install "${JAX_WHEEL}[tpu]" "$JAXLIB_WHEEL"
+            uv pip install $JAX_WHEEL[tpu] $JAXLIB_WHEEL
           fi
         env:
           INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
       - name: Verify JAX CUDA packages can be installed (Nvidia Pip Packages)
         run: |
-          # shellcheck disable=SC1090
-          "$PYTHON" -m uv venv ~/test_cuda_pip && source ~/test_cuda_pip/bin/activate
+          $PYTHON -m uv venv ~/test_cuda_pip && source ~/test_cuda_pip/bin/activate
           if [[ "${INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS}" == "1" ]]; then
-            uv pip install "${JAX_WHEEL}[cuda${{ matrix.cuda-version }}]"
+            uv pip install $JAX_WHEEL[cuda${{ matrix.cuda-version }}]
           else
-            uv pip install "${JAX_WHEEL}[cuda${{ matrix.cuda-version }}]" "$JAXLIB_WHEEL" "$JAX_CUDA_PJRT_WHEEL" "${JAX_CUDA_PLUGIN_WHEEL}[with-cuda]" "$EXTRA_INDEX"
+            uv pip install $JAX_WHEEL[cuda${{ matrix.cuda-version }}] $JAXLIB_WHEEL $JAX_CUDA_PJRT_WHEEL $JAX_CUDA_PLUGIN_WHEEL[with-cuda] $EXTRA_INDEX
           fi
         env:
           INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
       - name: Verify JAX CUDA packages can be installed (CUDA local)
         run: |
-          # ShellCheck cannot follow the 'activate' script because it is
-          # generated dynamically by the 'uv venv' command at runtime.
-          # shellcheck disable=SC1090
-          "$PYTHON" -m uv venv ~/test_cuda_local && source ~/test_cuda_local/bin/activate
+          $PYTHON -m uv venv ~/test_cuda_local && source ~/test_cuda_local/bin/activate
           if [[ "${INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS}" == "1" ]]; then
-            uv pip install "${JAX_WHEEL}[cuda${{ matrix.cuda-version }}-local]"
+            uv pip install $JAX_WHEEL[cuda${{ matrix.cuda-version }}-local]
           else
-            uv pip install "$JAX_WHEEL" "$JAXLIB_WHEEL" "$JAX_CUDA_PJRT_WHEEL" "$JAX_CUDA_PLUGIN_WHEEL"
+            uv pip install $JAX_WHEEL $JAXLIB_WHEEL $JAX_CUDA_PJRT_WHEEL $JAX_CUDA_PLUGIN_WHEEL
           fi
         env:
           INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
diff --git a/ci/run_pytest_rocm.sh b/ci/run_pytest_rocm.sh
index 1fe8582c1894..139542796256 100755
--- a/ci/run_pytest_rocm.sh
+++ b/ci/run_pytest_rocm.sh
@@ -110,6 +110,11 @@ export XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1 --xla_gpu_enable_ncc
 # Disable core dumps just in case
 ulimit -c 0
 
+# Prefer this checkout's `jax/` package over a jax preinstalled in site-packages
+# (e.g. CI images). Otherwise `import jax` can resolve to /usr/local/.../dist-packages
+# while `tests/` come from the PR tree, causing missing jax._src.* errors.
+export PYTHONPATH="$(pwd):${PYTHONPATH:-}"
+
 echo "Running ROCm tests..."
 export NPROC=32
 LOGS_DIR="logs"
diff --git a/conftest.py b/conftest.py
index 1e6e17ab966a..7c9d55325ca2 100644
--- a/conftest.py
+++ b/conftest.py
@@ -52,6 +52,10 @@ def add_imports(doctest_namespace):
 # For GPU, the env var JAX_ENABLE_CUDA_XDIST must be set equal to the number of
 # CUDA devices. Test processes will be assigned in round robin fashion across
 # the devices.
+#
+# For ROCm, use os.environ assignment (not setdefault) for ROCR_VISIBLE_DEVICES
+# so runner-injected visibility (e.g. GitHub gha-gpu-isolation-settings) cannot
+# defeat per-worker pinning. Device selection uses ROCR only (not HIP_VISIBLE_DEVICES).
 def pytest_collection() -> None:
   if os.environ.get("JAX_ENABLE_TPU_XDIST", None):
     # When running as an xdist worker, will be something like "gw0"
@@ -73,12 +77,20 @@ def pytest_collection() -> None:
         "CUDA_VISIBLE_DEVICES", str(xdist_worker_number % num_cuda_devices)
     )
 
-  elif num_rocm_devices := os.environ.get("JAX_ENABLE_ROCM_XDIST", None):
-    num_rocm_devices = int(num_rocm_devices)
+  elif os.environ.get("JAX_ENABLE_ROCM_XDIST", None) is not None:
     xdist_worker_name = os.environ.get("PYTEST_XDIST_WORKER", "")
     if not xdist_worker_name.startswith("gw"):
       return
-    xdist_worker_number = int(xdist_worker_name[len("gw") :])
-    os.environ.setdefault(
-        "ROCR_VISIBLE_DEVICES", str(xdist_worker_number % num_rocm_devices)
-    )
+    xdist_worker_number = int(xdist_worker_name[len("gw"):])
+
+    existing = os.environ.get("ROCR_VISIBLE_DEVICES", "")
+    if existing:
+      devices = existing.split(",")
+    else:
+      num = int(os.environ["JAX_ENABLE_ROCM_XDIST"])
+      devices = [str(i) for i in range(num)]
+
+    n = len(devices)
+    shift = xdist_worker_number % n
+    rotated = devices[shift:] + devices[:shift]
+    os.environ["ROCR_VISIBLE_DEVICES"] = ",".join(rotated)
diff --git a/tests/pallas/gpu_pallas_distributed_test.py b/tests/pallas/gpu_pallas_distributed_test.py
index e99a2e5cb2a4..a3e4c664e257 100644
--- a/tests/pallas/gpu_pallas_distributed_test.py
+++ b/tests/pallas/gpu_pallas_distributed_test.py
@@ -840,6 +840,8 @@ def _store():
 
 class PallasCallMultimemTest(TestCase):
   def setUp(self):
+    if cuda_versions is None:
+      self.skipTest("Multimem requires CUDA")
     if jax.device_count() < 2:
       self.skipTest("Needs at least two devices")
     if any(
@@ -1118,6 +1120,8 @@ class PallasCallMultimemThreadUnsafeTest(TestCase):
   """
 
   def setUp(self):
+    if cuda_versions is None:
+      self.skipTest("Multimem requires CUDA")
     if jax.local_device_count() > 1:
       self.skipTest("Multimem not supported in multi-thread mode yet.")
     if jax.device_count() < 2: