Test PR CI trigger

.github/actions/download-jax-rocm-wheels/action.yml

@@ -27,6 +27,9 @@ inputs:
   s3_download_uri:
     description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)"
     default: 'latest'
+  jax-jaxlib-source:
+    description: "Where to fetch the base jax/jaxlib wheels from: 'gcs' (upstream nightly bucket) or 's3' (built alongside the ROCm wheels on the fork's runners)"
+    default: 'gcs'
 permissions: {}
 runs:
   using: "composite"
@@ -74,7 +77,11 @@ runs:
         mkdir -p $(pwd)/dist
         # Disable Authentication on the ROCm runners.
         gcloud config set auth/disable_credentials True
-        if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then
+        if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then
+          # jax is fetched from S3/CloudFront in the jaxlib "head" block below,
+          # alongside jaxlib and the ROCm plugin/pjrt wheels.
+          echo "jax wheel will be downloaded from S3/CloudFront alongside the ROCm wheels."
+        elif [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then
           gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}"/jax*py3*none*any.whl $(pwd)/dist/
         else
           echo "JAX wheel won't be downloaded, only jaxlib pre-built wheel is tested."
@@ -86,7 +93,9 @@ runs:
           echo "JAX only release. Only downloading the jax wheel from the release bucket."
         else
           if [[ ${INPUTS_JAXLIB_VERSION} == "head" ]]; then
-            gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/
+            if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" != "s3" ]]; then
+              gcloud storage cp -r "${INPUTS_GCS_DOWNLOAD_URI}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/
+            fi
 
             # Resolve ROCm plugin/PJRT wheels path via CloudFront.
             ROCM_WHEELS_BASE_URL="https://d22q5eopkfeftw.cloudfront.net"
@@ -112,8 +121,31 @@ runs:
               exit 1
             fi
 
+            DOWNLOAD_URLS=("${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}")
+
+            # On the fork, the base jax/jaxlib wheels are built alongside the
+            # ROCm wheels and uploaded to the same S3 prefix, so fetch them from
+            # the same CloudFront listing instead of the (broken) GCS bucket.
+            if [[ "${INPUTS_JAX_JAXLIB_SOURCE}" == "s3" ]]; then
+              JAXLIB_FILE=$(echo "$FILES" | grep "jaxlib-" | grep "${PYTHON_MAJOR_MINOR}" | head -n1 || true)
+              if [[ -z "${JAXLIB_FILE}" ]]; then
+                echo "Could not find jaxlib wheel in ${WHEELS_URL}"
+                exit 1
+              fi
+              DOWNLOAD_URLS+=("${WHEELS_URL}/${JAXLIB_FILE}")
+
+              if [[ "${INPUTS_DOWNLOAD_JAX_FROM_GCS}" == "1" ]]; then
+                JAX_FILE=$(echo "$FILES" | grep -E "jax-[0-9].*-py3-none-any\.whl" | head -n1 || true)
+                if [[ -z "${JAX_FILE}" ]]; then
+                  echo "Could not find jax wheel in ${WHEELS_URL}"
+                  exit 1
+                fi
+                DOWNLOAD_URLS+=("${WHEELS_URL}/${JAX_FILE}")
+              fi
+            fi
+
             PYTHON_BIN="python${INPUTS_PYTHON}"
-            "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${WHEELS_URL}/${PJRT_FILE}" "${WHEELS_URL}/${PLUGIN_FILE}"
+            "${PYTHON_BIN}" -m pip download --no-deps --dest "$(pwd)/dist" "${DOWNLOAD_URLS[@]}"
           elif [[ ${INPUTS_JAXLIB_VERSION} == "pypi_latest" ]]; then
             PYTHON=python${INPUTS_PYTHON}
             $PYTHON -m pip download jaxlib jax-rocm${JAXCI_ROCM_VERSION}-pjrt jax-rocm${JAXCI_ROCM_VERSION}-plugin --dest $(pwd)/dist/ --no-deps
@@ -134,6 +166,7 @@ runs:
         INPUTS_PYTHON: ${{ inputs.python }}
         INPUTS_S3_DOWNLOAD_URI: ${{ inputs.s3_download_uri }}
         INPUTS_ROCM_FULL_VERSION: ${{ inputs.rocm-release-version || inputs.rocm-version }}
+        INPUTS_JAX_JAXLIB_SOURCE: ${{ inputs.jax-jaxlib-source }}
     - name: Skip the test run if the wheel artifacts were not downloaded successfully
       shell: bash
       if: steps.download-wheel-artifacts.outcome == 'failure'

.github/workflows/bazel_rocm.yml

@@ -87,6 +87,10 @@ on:
         description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)"
         default: 'latest'
         type: string
+      jax-jaxlib-source:
+        description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'"
+        default: 'gcs'
+        type: string
       gate-targets-file:
         description: "Path to a file listing Bazel targets to run (one per line). If empty, runs the full default test suite."
         type: string
@@ -181,6 +185,7 @@ jobs:
           jaxlib-version: ${{ inputs.jaxlib-version }}
           gcs_download_uri: ${{ inputs.gcs_download_uri }}
           s3_download_uri: ${{ inputs.s3_download_uri }}
+          jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }}
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c

.github/workflows/bazel_rocm_presubmit.yml

@@ -28,7 +28,6 @@ concurrency:
 permissions: {}
 jobs:
   run-tests:
-    if: github.event_name == 'workflow_dispatch' || (github.event.repository.fork == false && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm'))
     uses: ./.github/workflows/bazel_rocm.yml
 # Begin Presubmit Naming Check - name modification requires internal check to be updated
     strategy:
@@ -59,7 +58,6 @@ jobs:
       clone_main_xla: 0
 
   blocking-gate:
-    if: github.event_name == 'workflow_dispatch' || (github.event.repository.fork == false && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm'))
     uses: ./.github/workflows/bazel_rocm.yml
     strategy:
       fail-fast: false

.github/workflows/pytest_rocm.yml

@@ -66,6 +66,10 @@ on:
         description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)"
         type: string
         default: 'latest'
+      jax-jaxlib-source:
+        description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'"
+        type: string
+        default: 'gcs'
       halt-for-connection:
         description: 'Should this workflow run wait for a remote connection?'
         type: string
@@ -108,6 +112,10 @@ on:
         description: "S3 URI for ROCm plugin/PJRT wheels (use 'latest' to resolve via LATEST pointer)"
         default: 'latest'
         type: string
+      jax-jaxlib-source:
+        description: "Where to fetch base jax/jaxlib wheels from: 'gcs' or 's3'"
+        default: 'gcs'
+        type: string
       halt-for-connection:
         description: 'Should this workflow run wait for a remote connection?'
         type: string
@@ -168,6 +176,7 @@ jobs:
           skip-download-jaxlib-and-plugins-from-gcs: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
           gcs_download_uri: ${{ inputs.gcs_download_uri }}
           s3_download_uri: ${{ inputs.s3_download_uri }}
+          jax-jaxlib-source: ${{ inputs.jax-jaxlib-source }}
       - name: Install Python dependencies
         run: |
           $JAXCI_PYTHON -m pip install uv~=0.12.3

.github/workflows/scripts/select_continuous_wheel_test_jobs.py

@@ -39,6 +39,10 @@ class JobSpec:
     'build_jaxlib_artifact',
     'build_cuda_artifacts',
     'build_rocm_artifacts',
+    # Fork-only. Upstream builds jax/jaxlib in build_jax_artifact and
+    # build_jaxlib_artifact, which are gated on the jax-ml owner; on the ROCm
+    # fork this job builds them instead. See wheel_tests_continuous.yml.
+    'build_rocm_jax_jaxlib',
 )
 
 # requires_builds lists artifact jobs whose uploaded wheels are consumed by
@@ -118,6 +122,7 @@ class JobSpec:
             'build_jax_artifact',
             'build_jaxlib_artifact',
             'build_rocm_artifacts',
+            'build_rocm_jax_jaxlib',
         ),
     ),
     # LINT.ThenChange(.github/workflows/wheel_tests_continuous.yml:run_pytest_rocm)
@@ -130,6 +135,7 @@ class JobSpec:
             'build_jax_artifact',
             'build_jaxlib_artifact',
             'build_rocm_artifacts',
+            'build_rocm_jax_jaxlib',
         ),
     ),
     # LINT.ThenChange(.github/workflows/wheel_tests_continuous.yml:run_bazel_test_rocm)
@@ -145,6 +151,7 @@ class JobSpec:
 # callers keep halt-for-connection pinned to 'no'.
 ROCM_JOB_IDS = frozenset((
     'build_rocm_artifacts',
+    'build_rocm_jax_jaxlib',
     'run_pytest_rocm',
     'run_bazel_test_rocm',
 ))

.github/workflows/wheel_tests_continuous.yml

@@ -24,9 +24,12 @@
 #                           runs Bazel TPU tests with py_import.
 # 10. build-rocm-artifacts: Calls the `build_rocm_artifacts.yml` workflow to build ROCm plugin/pjrt
 #                           wheels and uploads them to an S3 bucket.
-# 11. run-pytest-rocm:      Calls the `pytest_rocm.yml` workflow which downloads the jaxlib and
+# 11. build-rocm-jax-jaxlib: Fork-only. Calls the `build_rocm_artifacts.yml` workflow to build the
+#                           base jax/jaxlib wheels into the same S3 prefix, since the upstream
+#                           jax/jaxlib build jobs above only run under the jax-ml owner.
+# 12. run-pytest-rocm:      Calls the `pytest_rocm.yml` workflow which downloads the jaxlib and
 #                           ROCm artifacts and runs the ROCm tests.
-# 12. run-bazel-test-rocm:  Calls the `bazel_rocm.yml` workflow which runs the ROCm Bazel tests.
+# 13. run-bazel-test-rocm:  Calls the `bazel_rocm.yml` workflow which runs the ROCm Bazel tests.
 
 name: CI - Wheel Tests (Continuous)
 permissions:
@@ -524,9 +527,51 @@ jobs:
       upload_artifacts_to_s3: true
       s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}'
 
+  # build-jax-artifact and build-jaxlib-artifact above are gated on the jax-ml
+  # owner, so on the fork the ROCm test legs would fall back to the upstream
+  # nightly GCS bucket, which is frequently stale or missing. Build the base
+  # jax and jaxlib wheels here instead, on the same runners and into the same
+  # S3 prefix as the ROCm plugin/pjrt wheels, mirroring what
+  # wheel_tests_nightly_release.yml already does.
+  build-rocm-jax-jaxlib:
+    needs: [select-jobs]
+    if: >-
+      ${{
+          fromJSON(needs.select-jobs.outputs.required_builds || '{}').build_rocm_jax_jaxlib == true
+          && github.repository_owner == 'ROCm'
+      }}
+    uses: ./.github/workflows/build_rocm_artifacts.yml
+    permissions:
+      id-token: write
+      contents: read
+      actions: read
+    strategy:
+        fail-fast: false
+        matrix:
+          runner: ["amd-do-linux.jax.cpu"]
+          artifact: ["jax", "jaxlib"]
+          # Only the Python versions consumed by the ROCm test legs below are
+          # built here; unlike the plugin/pjrt wheels these are not published.
+          python: ["3.12"]
+          rocm: [
+            {label: "TheRock 10.0.0", wheel-version: "10", release-version: "10.0.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-10.0:latest" },
+          ]
+    name: "Build ROCm jax/jaxlib ${{ matrix.rocm.label }}"
+    with:
+      runner: ${{ matrix.runner }}
+      artifact: ${{ matrix.artifact }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm.wheel-version }}
+      rocm-release-version: ${{ matrix.rocm.release-version }}
+      manylinux-image: ${{ matrix.rocm.manylinux-image }}
+      clone_main_xla: 0
+      halt-for-connection: 'no'
+      upload_artifacts_to_s3: true
+      s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}'
+
   # LINT.IfChange(run_pytest_rocm)
   run-pytest-rocm:
-    needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts]
+    needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts, build-rocm-jax-jaxlib]
   # LINT.ThenChange(.github/workflows/scripts/select_continuous_wheel_test_jobs.py:run_pytest_rocm)
     if: >-
       ${{
@@ -557,11 +602,14 @@ jobs:
       jaxlib-version: "head"
       gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri || 'gs://jax-nightly-artifacts/latest' }}
       s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}'
+      # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of
+      # the upstream GCS nightly bucket, which is frequently stale/broken.
+      jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }}
       halt-for-connection: 'no'
 
   # LINT.IfChange(run_bazel_test_rocm)
   run-bazel-test-rocm:
-    needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts]
+    needs: [select-jobs, build-jax-artifact, build-jaxlib-artifact, build-rocm-artifacts, build-rocm-jax-jaxlib]
   # LINT.ThenChange(.github/workflows/scripts/select_continuous_wheel_test_jobs.py:run_bazel_test_rocm)
     if: >-
       ${{
@@ -593,6 +641,9 @@ jobs:
       enable-x64: ${{ matrix.enable-x64 }}
       gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri || 'gs://jax-nightly-artifacts/latest' }}
       s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-continuous/${{ github.run_number }}/${{ github.run_attempt }}'
+      # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of
+      # the upstream GCS nightly bucket, which is frequently stale/broken.
+      jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }}
       build_jaxlib: ${{ 'false' }}
       build_jax: ${{ 'false' }}
       jaxlib-version: "head"

.github/workflows/wheel_tests_nightly_release.yml

@@ -221,9 +221,49 @@ jobs:
       s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}'
       git_commit: ${{ inputs.git_commit }}
 
+  # Upstream CI frequently stops publishing jax/jaxlib nightly wheels to GCS.
+  # On the ROCm fork, build the base jax and jaxlib wheels ourselves on the same
+  # runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix,
+  # so the ROCm test legs can consume them instead of the (stale) GCS bucket.
+  build-rocm-jax-jaxlib:
+    if: github.repository_owner == 'ROCm'
+    uses: ./.github/workflows/build_rocm_artifacts.yml
+    permissions:
+      id-token: write
+      contents: read
+      actions: read
+    strategy:
+        fail-fast: false
+        matrix:
+          runner: ["amd-do-linux.jax.cpu"]
+          artifact: ["jax", "jaxlib"]
+          python: ["3.12", "3.13", "3.14"]
+          rocm: [
+            {label: "TheRock 10.0.0", wheel-version: "10", release-version: "10.0.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-10.0:latest" },
+            {label: "TheRock latest", wheel-version: "10", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"},
+          ]
+          exclude:
+            # jax is a pure-python (py3-none-any) wheel; build it once per ROCm
+            # variant rather than once per Python version.
+            - artifact: "jax"
+              python: "3.13"
+            - artifact: "jax"
+              python: "3.14"
+    name: "Build ROCm jax/jaxlib (${{ matrix.rocm.label }})"
+    with:
+      runner: ${{ matrix.runner }}
+      artifact: ${{ matrix.artifact }}
+      python: ${{ matrix.python }}
+      rocm-version: ${{ matrix.rocm.wheel-version }}
+      rocm-release-version: ${{ matrix.rocm.release-version }}
+      manylinux-image: ${{ matrix.rocm.manylinux-image }}
+      clone_main_xla: 0
+      upload_artifacts_to_s3: true
+      s3_upload_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}'
+
   run-pytest-rocm:
     if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }}
-    needs: [build-rocm-artifacts]
+    needs: [build-rocm-artifacts, build-rocm-jax-jaxlib]
     permissions:
       id-token: write
       contents: read
@@ -250,10 +290,13 @@ jobs:
       gcs_download_uri: ${{inputs.gcs_download_uri}}
       s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}'
       git_commit: ${{ inputs.git_commit }}
+      # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of
+      # the upstream GCS nightly bucket, which is frequently stale/broken.
+      jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }}
 
   run-bazel-test-rocm:
     if: ${{ !cancelled() && (github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm') }}
-    needs: [build-rocm-artifacts]
+    needs: [build-rocm-artifacts, build-rocm-jax-jaxlib]
     permissions:
       id-token: write
       contents: read
@@ -279,6 +322,9 @@ jobs:
       gcs_download_uri: ${{inputs.gcs_download_uri}}
       halt-for-connection: ${{inputs.halt-for-connection}}
       s3_download_uri: 's3://jax-ci-amd/rocm-wheels/${{ github.repository }}/${{ github.ref_name }}/${{ matrix.rocm.release-version }}/wheel-tests-nightly/${{ github.run_number }}/${{ github.run_attempt }}'
+      # On the ROCm fork, use the jax/jaxlib wheels built above (S3) instead of
+      # the upstream GCS nightly bucket, which is frequently stale/broken.
+      jax-jaxlib-source: ${{ github.repository_owner == 'ROCm' && 's3' || 'gcs' }}
       build_jaxlib: ${{ 'false' }}
       build_jax: ${{ 'false' }}
       jaxlib-version: "head"

build/rocm/ci_test_targets.txt

@@ -7,50 +7,8 @@
 //tests/pallas:gpu_tests
 //tests/pallas:backend_independent_tests
 //jaxlib/tools:check_gpu_wheel_sources_test
--//tests/pallas:pallas_test_gpu
--//tests/pallas:ops_test_gpu
+
+# Mosaic GPU (mgpu) ops test. Mosaic GPU is an NVIDIA-only codegen path, so
+# this never passes on ROCm. Note it carries no multiaccelerator tag, so this
+# exclusion is what keeps it out of single_gpu runs -- not the tag filter.
 -//tests/pallas:ops_test_mgpu_gpu
--//tests/pallas:pallas_shape_poly_test_gpu
--//tests/pallas:pallas_vmap_test_gpu
--//tests/pallas:triton_pallas_test_gpu
--//tests/pallas:gpu_paged_attention_test_gpu
--//tests:export_harnesses_multi_platform_test_gpu
--//tests:jet_test_gpu
--//tests:lax_autodiff_test_gpu
--//tests:lax_numpy_setops_test_gpu
--//tests:lax_numpy_test_gpu
--//tests:lax_test_gpu
--//tests:linalg_test_gpu
--//tests:logging_test_gpu
--//tests:random_lax_test_gpu
--//tests:scipy_signal_test_gpu
--//tests:stax_test_gpu
--//tests:ode_test_gpu
--//tests:lobpcg_test_gpu
--//tests:scipy_stats_test_gpu
--//tests:nn_test_gpu
--//tests:lax_scipy_sparse_test_gpu
--//tests:lax_scipy_spectral_dac_test_gpu
--//tests:lax_scipy_special_functions_test_gpu
--//tests:cholesky_update_test_gpu
--//tests:api_test_gpu
--//tests:ann_test_gpu
--//tests:experimental_rnn_test_gpu
--//tests:lax_vmap_test_gpu
--//tests:qdwh_test_gpu
--//tests:scaled_dot_test_gpu
--//tests:scipy_spatial_test_gpu
--//tests:shape_poly_test_gpu
--//tests:sparsify_test_gpu
--//tests:lax_numpy_reducers_test_gpu
--//tests:scipy_optimize_test_gpu
--//tests:lax_numpy_einsum_test_gpu
--//tests:lax_scipy_test_gpu
--//tests:sparse_test_gpu
--//tests:sparse_bcoo_bcsr_test_gpu
--//tests:svd_test_gpu
--//tests:eigh_test_gpu
--//tests:image_test_gpu
-# buffer_callback_test_gpu uses py_import deps that are incompatible with
-# pre-built plugin wheels (build_jaxlib=false); exclude it unconditionally.
--//tests:buffer_callback_test_gpu