Rocm jaxlib v0.9.0

.bazelversion

@@ -1 +1 @@
-7.4.1
+7.7.0

.editorconfig

@@ -5,18 +5,7 @@ indent_style = space
 end_of_line = lf
 trim_trailing_whitespace = true
 insert_final_newline = true
-
-[*.py]
-max_line_length = 79
-indent_size = 2
-
-[*.rst]
-max_line_length = 79
 indent_size = 2
 
-[*.md]
-max_line_length = 79
-indent_size = 2
-
-[*.yml]
-indent_size = 2
+[*.py]
+max_line_length = 80

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -24,7 +24,7 @@ body:
 
       [issue search]: https://github.com/jax-ml/jax/search?q=is%3Aissue&type=issues
 
-      [Raw report]: http://github.com/jax-ml/jax/issues/new
+      [Raw report]: https://github.com/jax-ml/jax/issues/new?template=none
 - type: textarea
   attributes:
     label: Description

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,11 @@
+<!--
+
+Thanks for taking the time to contribute to JAX! A couple things to keep in mind:
+
+* Contributing to JAX requires signing the Contributor License Agreement: https://docs.jax.dev/en/latest/contributing.html#google-contributor-license-agreement
+
+* Please run lint checks and tests locally: https://docs.jax.dev/en/latest/contributing.html#contributing-code-using-pull-requests
+
+* If applicable, read our policy on AI generated code: https://docs.jax.dev/en/latest/contributing.html#can-i-contribute-ai-generated-code
+
+-->

.github/actionlint.yaml

@@ -0,0 +1,20 @@
+# Configuration related to self-hosted runner.
+self-hosted-runner:
+  labels:
+    - "linux-x86-n4-16" # Linux X86 runner using the 16 vcpu n4-standard-16 machine.
+    - "linux-x86-n4-32" # Linux X86 runner using the 32 vcpu n4-standard-32 machine.
+    - "linux-x86-n4-64" # Linux X86 runner using the 64 vcpu n2-standard-64 machine.
+    - "linux-x86-g2-16-l4-1gpu" # Linux X86 GPU runner using g2-standard-16 machine with 1 NVIDIA L4 GPU attached.
+    - "linux-x86-g2-48-l4-4gpu" # Linux X86 GPU runner using g2-standard-48 machine with 4 NVIDIA L4 GPUs attached.
+    - "linux-x86-ct5lp-224-8tpu" # Linux X86 TPU runner using ct5lp-hightpu-8t machine with 2x4 topology.
+    - "linux-arm64-c4a-16" # Linux ARM64 CPU Runner using the 16 vcpu c4a-standard-16 machine.
+    - "linux-arm64-c4a-64" # Linux ARM64 CPU Runner using the 64 vcpu c4a-standard-64 machine.
+    - "windows-x86-n2-16" # Windows X86 runner using n2-standard-16 machine.
+    - "windows-x86-n2-64" # Windows X86 runner using n2-standard-64 machine.
+    - "linux-x86-a4-224-b200-1gpu" # Linux X86 GPU runner using 1 B200 GPU and 1/8 the resources of a a4-highgpu-8g machine
+    - "linux-x86-a3-8g-h100-8gpu" # Linux X86 GPU runner using a3-highgpu-8g machine with 8 NVIDIA H100 GPUs attached.
+    - "linux-x86-ct6e-180-8tpu" # Linux X86 TPU runner using ct6e-hightpu-8t machine with 2x4 topology.
+    - "linux-x86-ct6e-180-4tpu" # Linux X86 TPU runner using ct6e-hightpu-4t machine with 2x2 topology.
+    - "linux-x86-ct4p-240-4tpu" # Linux X86 TPU runner using ct4p-hightpu-4t machine with 2x2x1 topology.
+    - "linux-x86-tpu7x-224-4tpu" # Linux X86 TPU runner using tpu7x-224 machine with 4 TPU chips (8 cores) and 2x2x1 topology.
+    - "linux-x86_64-cirrascale-64-8gpu-amd-mi250" # AMD runner

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

@@ -0,0 +1,105 @@
+# Composite action to download the jax and jaxlib wheels
+name: Download JAX CPU wheels
+
+inputs:
+  runner:
+    description: "Which runner type should the wheels be downloaded for?"
+    type: string
+    default: "linux-x86-n4-16"
+  python:
+    description: "Which python version should the artifact be downloaded for?"
+    required: true
+    type: string
+  jaxlib-version:
+    description: "Which jaxlib version to download? (head/pypi_latest)"
+    type: string
+    default: "head"
+  skip-download-jaxlib-from-gcs:
+    description: "Whether to skip downloading the jaxlib artifact from GCS (e.g for testing a jax only release)"
+    default: '0'
+    type: string
+  gcs_download_uri:
+    description: "GCS location prefix from where the artifacts should be downloaded"
+    default: 'gs://general-ml-ci-transient/jax-github-actions/jax/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
+    type: string
+permissions: {}
+runs:
+  using: "composite"
+
+  steps:
+    # Note that certain envs such as JAXCI_HERMETIC_PYTHON_VERSION are set by the calling workflow.
+    - name: Set env vars for use in artifact download URL
+      shell: bash
+      run: |
+        os=$(uname -s | awk '{print tolower($0)}')
+        arch=$(uname -m)
+
+        # Adjust os and arch for Windows
+        if [[  $os  =~ "msys_nt" ]] && [[ $arch =~ "x86_64" ]]; then
+          os="win"
+          arch="amd64"
+        fi
+
+        # Get the major and minor version of Python.
+        # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.10, then python_major_minor=310
+        # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.13-nogil, then python_major_minor=313t
+        python_major_minor=$(echo "${JAXCI_HERMETIC_PYTHON_VERSION//-nogil/t}" | tr -d '.')
+
+        echo "OS=${os}" >> $GITHUB_ENV
+        echo "ARCH=${arch}" >> $GITHUB_ENV
+        # Python wheels follow a naming convention: standard wheels use the pattern
+        # `*-cp<py_version>-cp<py_version>-*`, while free-threaded wheels use
+        # `*-cp<py_version>-cp<py_version>t-*`.
+        echo "PYTHON_MAJOR_MINOR=cp${python_major_minor%t}-cp${python_major_minor}-" >> $GITHUB_ENV
+    - name: Download wheels from GCS (non-Windows runs)
+      shell: bash
+      id: download-wheel-artifacts-nw
+      # Set continue-on-error to true to prevent actions from failing the workflow if this step
+      # fails. Instead, we verify the outcome in the step below so that we can print a more
+      # informative error message.
+      continue-on-error: true
+      if: ${{ !contains(inputs.runner, 'windows-x86') }}
+      run: |
+        mkdir -p $(pwd)/dist
+        gcloud storage cp -r "${{ inputs.gcs_download_uri }}"/jax*py3*none*any.whl $(pwd)/dist/
+
+        if [[ "${{ inputs.skip-download-jaxlib-from-gcs }}" == "1" ]]; then
+          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/
+          elif [[ ${{ inputs.jaxlib-version }} == "pypi_latest" ]]; then
+            PYTHON=python${{ inputs.python }}
+            $PYTHON -m pip download jaxlib --dest $(pwd)/dist/
+          else
+            echo "Invalid jaxlib version: ${{ inputs.jaxlib-version }}"
+            exit 1
+          fi
+        fi
+    - name: Download wheels from GCS (Windows runs)
+      shell: cmd
+      id: download-wheel-artifacts-w
+      # Set continue-on-error to true to prevent actions from failing the workflow if this step
+      # fails. Instead, we verify the outcome in step below so that we can print a more
+      # informative error message.
+      continue-on-error: true
+      if: ${{ contains(inputs.runner, 'windows-x86') }}
+      run: |
+        mkdir dist
+        @REM Use `call` so that we can run sequential gcloud storage commands on Windows
+        @REM See https://github.com/GoogleCloudPlatform/gsutil/issues/233#issuecomment-196150652
+        call gcloud storage cp -r "${{ inputs.gcs_download_uri }}"/jax*py3*none*any.whl dist/
+
+        if "${{ inputs.skip-download-jaxlib-from-gcs }}"=="1" (
+          echo "JAX only release. Only downloading the jax wheel from the release bucket."
+        ) else (
+          call gcloud storage cp -r "${{ inputs.gcs_download_uri }}/jaxlib*%PYTHON_MAJOR_MINOR%*%OS%*%ARCH%*.whl" dist/
+        )
+    - name: Skip the test run if the wheel artifacts were not downloaded successfully
+      shell: bash
+      if: steps.download-wheel-artifacts-nw.outcome == 'failure' || steps.download-wheel-artifacts-w.outcome == 'failure'
+      run: |
+        echo "Failed to download wheel artifacts from GCS. Please check if the wheels were"
+        echo "built successfully by the artifact build jobs and are available in the GCS bucket."
+        echo "Skipping the test run."
+        exit 1

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

@@ -0,0 +1,100 @@
+# Composite action to download the jax, jaxlib, and the CUDA plugin wheels
+name: Download JAX CUDA wheels
+
+inputs:
+  python:
+    description: "Which python version should the artifact be downloaded for?"
+    type: string
+    required: true
+  cuda-version:
+    description: "Which cuda version should the artifact be downloaded for?"
+    type: string
+    default: "12"
+  use-nvidia-pip-wheels:
+    description: "Whether to download Nvidia CUDA packages from PyPI?"
+    type: boolean
+    default: false
+  jaxlib-version:
+    description: "Which jaxlib version to download? (head/pypi_latest)"
+    type: string
+    default: "head"
+  download-jax-from-gcs:
+    description: "Whether to download the jax wheel from GCS"
+    default: '1'
+    type: string
+  skip-download-jaxlib-and-cuda-plugins-from-gcs:
+    description: "Whether to skip downloading the jaxlib and cuda plugins from GCS (e.g for testing a jax only release)"
+    default: '0'
+    type: string
+  gcs_download_uri:
+    description: "GCS location prefix from where the artifacts should be downloaded"
+    default: 'gs://general-ml-ci-transient/jax-github-actions/jax/${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}'
+    type: string
+permissions: {}
+runs:
+  using: "composite"
+
+  steps:
+    # Note that certain envs such as JAXCI_HERMETIC_PYTHON_VERSION are set by the calling workflow.
+    - name: Set env vars for use in artifact download URL
+      shell: bash
+      run: |
+        os=$(uname -s | awk '{print tolower($0)}')
+        arch=$(uname -m)
+
+        # Get the major and minor version of Python.
+        # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.11, then python_major_minor=311
+        # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.13-nogil, then python_major_minor=313t
+        python_major_minor=$(echo "${JAXCI_HERMETIC_PYTHON_VERSION//-nogil/t}" | tr -d '.')
+
+        echo "OS=${os}" >> $GITHUB_ENV
+        echo "ARCH=${arch}" >> $GITHUB_ENV
+        # Python wheels follow a naming convention: standard wheels use the pattern
+        # `*-cp<py_version>-cp<py_version>-*`, while free-threaded wheels use
+        # `*-cp<py_version>-cp<py_version>t-*`.
+        echo "PYTHON_MAJOR_MINOR=cp${python_major_minor%t}-cp${python_major_minor}-" >> $GITHUB_ENV
+
+        # Get the CUDA major version only
+        full_cuda_version="${{ inputs.cuda-version }}"
+        echo "JAXCI_CUDA_VERSION=${full_cuda_version%%.*}" >> $GITHUB_ENV
+    - name: Download wheels
+      shell: bash
+      id: download-wheel-artifacts
+      # Set continue-on-error to true to prevent actions from failing the workflow if this step
+      # fails. Instead, we verify the outcome in the next step so that we can print a more
+      # informative error message.
+      continue-on-error: true
+      run: |
+        mkdir -p $(pwd)/dist
+        if [[ "${{ 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."
+        fi
+
+        # Do not download the jaxlib and CUDA plugin artifacts if we are testing a jax only
+        # release.
+        if [[ "${{ inputs.skip-download-jaxlib-and-cuda-plugins-from-gcs }}" == "1" ]]; then
+          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/
+            gcloud storage cp -r "${{ inputs.gcs_download_uri }}/jax*cuda${JAXCI_CUDA_VERSION}*plugin*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/
+            gcloud storage cp -r "${{ inputs.gcs_download_uri }}/jax*cuda${JAXCI_CUDA_VERSION}*pjrt*${OS}*${ARCH}*.whl" $(pwd)/dist/
+          elif [[ ${{ inputs.jaxlib-version }} == "pypi_latest" ]]; then
+            PYTHON=python${{ inputs.python }}
+            $PYTHON -m pip download jaxlib jax-cuda${JAXCI_CUDA_VERSION}-pjrt jax-cuda${JAXCI_CUDA_VERSION}-plugin --dest $(pwd)/dist/
+          else
+            echo "Invalid jaxlib version: ${{ inputs.jaxlib-version }}"
+            exit 1
+          fi
+        fi
+    - name: Skip the test run if the wheel artifacts were not downloaded successfully
+      shell: bash
+      if: steps.download-wheel-artifacts.outcome == 'failure'
+      run: |
+        echo "Failed to download wheel artifacts. Please check if the wheels were"
+        echo "built successfully by the artifact build jobs and are available in the GCS bucket if
+        echo "downloading from GCS."
+        echo "Skipping the test run."
+        exit 1

.github/workflows/asan.yaml

@@ -13,12 +13,17 @@ on:
       - main
     paths:
       - '**/workflows/asan.yaml'
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
 
 jobs:
   asan:
     # Don't execute in fork due to runner type
     if: github.repository == 'jax-ml/jax'
-    runs-on: linux-x86-n2-64
+    runs-on: linux-x86-n4-64
     container:
       image: index.docker.io/library/ubuntu@sha256:b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b # ratchet:ubuntu:24.04
     strategy:
@@ -38,14 +43,16 @@ jobs:
             zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev curl git \
             libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev \
             libffi-dev liblzma-dev
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
         with:
           path: jax
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+          persist-credentials: false
+      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
         with:
           repository: python/cpython
           path: cpython
           ref: v3.13.0
+          persist-credentials: false
       - name: Build CPython with ASAN enabled
         env:
           ASAN_OPTIONS: detect_leaks=0
@@ -60,6 +67,7 @@ jobs:
         env:
           ASAN_OPTIONS: detect_leaks=0
         run: |
+          apt install -y xxd
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
           pip install uv~=0.5.30
@@ -72,8 +80,7 @@ jobs:
           cd jax
           python build/build.py build --wheels=jaxlib --verbose \
             --bazel_options=--color=yes \
-            --bazel_options=--copt=-fsanitize=address \
-            --clang_path=/usr/bin/clang-18
+            --bazel_options=--copt=-fsanitize=address
           uv pip install dist/jaxlib-*.whl \
             -e .
       - name: Run tests