[Test do not merge] Testing build wheels

.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,113 @@
+# 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
+      env:
+        INPUTS_GCS_DOWNLOAD_URI: ${{ inputs.gcs_download_uri }}
+        INPUTS_SKIP_DOWNLOAD_JAXLIB_FROM_GCS: ${{ inputs.skip-download-jaxlib-from-gcs }}
+        INPUTS_JAXLIB_VERSION: ${{ inputs.jaxlib-version }}
+        INPUTS_PYTHON: ${{ inputs.python }}
+    - 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/
+        )
+      env:
+        INPUTS_GCS_DOWNLOAD_URI: ${{ inputs.gcs_download_uri }}
+        INPUTS_SKIP_DOWNLOAD_JAXLIB_FROM_GCS: ${{ inputs.skip-download-jaxlib-from-gcs }}
+    - 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,108 @@
+# 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-plugins-from-gcs:
+    description: "Whether to skip downloading the jaxlib and 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
+      env:
+        INPUTS_CUDA_VERSION: ${{ inputs.cuda-version }}
+    - 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_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
+      env:
+        INPUTS_DOWNLOAD_JAX_FROM_GCS: ${{ inputs.download-jax-from-gcs }}
+        INPUTS_GCS_DOWNLOAD_URI: ${{ inputs.gcs_download_uri }}
+        INPUTS_SKIP_DOWNLOAD_JAXLIB_AND_PLUGINS_FROM_GCS: ${{ inputs.skip-download-jaxlib-and-plugins-from-gcs }}
+        INPUTS_JAXLIB_VERSION: ${{ inputs.jaxlib-version }}
+        INPUTS_PYTHON: ${{ inputs.python }}
+    - 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