TheRock/.github/workflows/multi_arch_release_linux_jax_wheels.yml

# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

# Multi-arch release entry point for Linux JAX wheels.
#
# This workflow owns the release matrix and delegates each matrix cell to the
# reusable single-build workflow.

name: Multi-Arch Release Linux JAX Wheels

on:
  workflow_call:
    inputs:
      test_amdgpu_family:
        description: "GPU family to use for JAX tests; empty means skip tests."
        type: string
        default: ""
      release_type:
        description: 'Release type: "dev", "nightly", or "prerelease".'
        type: string
        required: true
      rocm_version:
        description: "ROCm package version to build against."
        type: string
        required: true
      rocm_package_find_links_url:
        description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build."
        type: string
        default: ""
      repository:
        description: "Repository to checkout. Defaults to github.repository."
        type: string
        default: ""
      ref:
        description: "Branch, tag, or SHA to checkout."
        type: string
        default: ""
      python_version:
        description: "Single Python version to build (empty for full matrix)"
        type: string
        default: ""
  workflow_dispatch:
    inputs:
      test_amdgpu_family:
        description: "Representative GPU family to use for JAX tests."
        type: string
        default: "gfx94X-dcgpu"
      release_type:
        description: 'Release type; developer-triggered jobs should use "dev".'
        type: choice
        options:
          - dev
        default: dev
      rocm_version:
        description: "ROCm package version to build against."
        type: string
        required: true
      rocm_package_find_links_url:
        description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build."
        type: string
        default: ""
      repository:
        description: "Repository to checkout. Defaults to github.repository."
        type: string
        default: ""
      ref:
        description: "Branch, tag, or SHA to checkout."
        type: string
        default: ""
      python_version:
        description: "Single Python version to build (empty for full matrix)"
        type: string
        default: ""

run-name: Release Linux JAX Wheels (${{ inputs.release_type }}, ${{ inputs.rocm_version }}, Multiarch)

permissions:
  contents: read
jobs:
  setup_matrix:
    runs-on: ubuntu-24.04
    outputs:
      jax_matrix: ${{ steps.matrix.outputs.jax_matrix }}
    steps:
      - name: Checkout
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: ${{ inputs.repository || github.repository }}
          ref: ${{ inputs.ref }}

      - name: Generate JAX matrix
        id: matrix
        run: |
          python ./build_tools/github_actions/configure_jax_release_matrix.py \
            --python-versions="${{ inputs.python_version }}"

  build_jax_wheels:
    needs: setup_matrix
    name: Build | py ${{ matrix.python_version }} | jax ${{ matrix.jax_ref }}
    strategy:
      fail-fast: false
      matrix:
        include: ${{ fromJSON(needs.setup_matrix.outputs.jax_matrix) }}

    permissions:
      id-token: write
      contents: read
    uses: ./.github/workflows/multi_arch_build_linux_jax_wheels.yml
    with:
      test_amdgpu_family: ${{ inputs.test_amdgpu_family }}
      python_version: ${{ matrix.python_version }}
      jax_ref: ${{ matrix.jax_ref }}
      jax_repository: ${{ matrix.jax_repository }}
      gfx_arch: ${{ matrix.gfx_arch }}
      rocm_version: ${{ inputs.rocm_version }}
      rocm_package_find_links_url: ${{ inputs.rocm_package_find_links_url }}
      release_type: ${{ inputs.release_type }}
      repository: ${{ inputs.repository || github.repository }}
      ref: ${{ inputs.ref }}