jax-array-api.yml

name: JAX Array API

on:
  workflow_dispatch: # allows triggering the workflow run manually
  pull_request:  # Automatically trigger on pull requests affecting particular files
    branches:
      - main
    paths:
      - '**workflows/jax-array-api.yml'
      - '**experimental/array_api/**'

jobs:
  build:

    runs-on: ubuntu-latest
    strategy:
      matrix:
        python-version: [3.11]

    steps:
    - name: Checkout jax
      uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet: actions/checkout@v4
    - name: Checkout array-api-tests
      uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet: actions/checkout@v4
      with:
        repository: data-apis/array-api-tests
        # TODO(jakevdp) update this to a stable release/tag when available.
        ref: 'db95e67b29235249e5776ca2b6bb4e77117e0690'  # Latest commit as of 2024-08-08
        submodules: 'true'
        path: 'array-api-tests'
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # ratchet:actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
    - name: Install dependencies
      run: |
        python -m pip install .[ci]
        python -m pip install -r array-api-tests/requirements.txt
    - name: Run the test suite
      env:
        ARRAY_API_TESTS_MODULE: jax.numpy
        JAX_ENABLE_X64: 'true'
      run: |
        cd ${GITHUB_WORKSPACE}/array-api-tests
        pytest array_api_tests --max-examples=5 --derandomize --disable-deadline --skips-file ${GITHUB_WORKSPACE}/tests/array_api_skips.txt