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

name: Test JAX Wheels

on:
  workflow_dispatch:
    inputs:
      test_runs_on:
        required: true
        type: string
        default: "linux-gfx942-1gpu-ccs-csp-ossci-rocm"
      image_name:
        required: true
        description: JAX docker image to run tests with
        type: string
      jax_version:
        description: Version of JAX to install
        required: false
        type: string
      jax_plugin_branch:
        required: true
        description: JAX plugin branch to checkout
        type: string
        default: "rocm-jaxlib-v0.11.0"

  workflow_call:
    inputs:
      test_runs_on:
        required: true
        type: string
      image_name:
        required: true
        description: JAX docker image to run tests with
        type: string
      jax_version:
        description: Version of JAX to install instead of the one on the docker image
        required: false
        type: string
      jax_plugin_branch:
        description: JAX plugin branch to checkout to use for test scripts
        type: string
        default: "rocm-jaxlib-v0.11.0"

permissions:
  contents: read

jobs:
  test_wheels:
    name: Test
    runs-on: ${{ inputs.test_runs_on }}
    steps:
      - name: Checkout
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repo: rocm/rocm-jax
      # TODO: Add steps for creating the JAX docker image with an install of TheRock and then running JAX tests on the container