rocm_jax/.github/workflows/ci-build.yaml

name: CI

# We test all supported Python versions as follows:
# - 3.7  : Documentation build
# - 3.8  : Part of Matrix
# - 3.9  : Part of Matrix with NumPy dispatch
# - 3.10 : Part of Matrix

on:
  # Trigger the workflow on push or pull request,
  # but only for the main branch
  push:
    branches:
      - main
  pull_request:
    branches:
      - main

jobs:
  lint_and_typecheck:
    runs-on: ubuntu-latest
    timeout-minutes: 5
    steps:
      - name: Cancel previous
        uses: styfle/cancel-workflow-action@0.9.1
        with:
          access_token: ${{ github.token }}
        if: ${{github.ref != 'refs/head/main'}}
      - uses: actions/checkout@v3
      - name: Set up Python 3.8
        uses: actions/setup-python@v3
        with:
          python-version: 3.8
      - uses: pre-commit/action@v2.0.3

  build:
    name: "build ${{ matrix.name-prefix }} (py ${{ matrix.python-version }} on ${{ matrix.os }}, x64=${{ matrix.enable-x64}})"
    runs-on: ${{ matrix.os }}
    timeout-minutes: 60
    strategy:
      matrix:
        include:
          - name-prefix: "with 3.8"
            python-version: "3.8"
            os: ubuntu-latest
            enable-x64: 0
            package-overrides: "none"
            num_generated_cases: 1
            use-latest-jaxlib: false
          - name-prefix: "with numpy-dispatch"
            python-version: "3.9"
            os: ubuntu-latest
            enable-x64: 1
            # Test experimental NumPy dispatch
            package-overrides: "git+https://github.com/seberg/numpy-dispatch.git"
            num_generated_cases: 1
            use-latest-jaxlib: false
          - name-prefix: "with 3.10"
            python-version: "3.10"
            os: ubuntu-latest
            enable-x64: 0
            package-overrides: "none"
            num_generated_cases: 1
            use-latest-jaxlib: false
    steps:
    - name: Cancel previous
      uses: styfle/cancel-workflow-action@0.9.1
      with:
        access_token: ${{ github.token }}
      if: ${{github.ref != 'refs/head/main'}}
    - uses: actions/checkout@v3
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v3
      with:
        python-version: ${{ matrix.python-version }}
    - name: Get pip cache dir
      id: pip-cache
      run: |
        python -m pip install --upgrade pip wheel
        echo "::set-output name=dir::$(pip cache dir)"
    - name: pip cache
      uses: actions/cache@v3
      with:
        path: ${{ steps.pip-cache.outputs.dir }}
        key: ${{ runner.os }}-pip-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt') }}
        restore-keys: |
          ${{ runner.os }}-pip-
    - name: Install dependencies
      run: |
        pip install -r build/test-requirements.txt
        if [ "${{ matrix.package-overrides }}" != "none" ]; then
          pip install ${{ matrix.package-overrides }}
        fi
        if [ "${{ matrix.use-latest-jaxlib }}" == "true" ]; then
          pip install .[cpu]
        else
          pip install .[minimum-jaxlib]
        fi

    - name: Run tests
      env:
        JAX_NUM_GENERATED_CASES: ${{ matrix.num_generated_cases }}
        JAX_ENABLE_X64: ${{ matrix.enable-x64 }}
        JAX_ENABLE_CHECKS: true
      run: |
        pip install -e .
        echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
        echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
        echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
        pytest -n auto --tb=short tests examples


  documentation:
    runs-on: ubuntu-latest
    timeout-minutes: 5
    strategy:
      matrix:
        python-version: [3.7]
    steps:
    - name: Cancel previous
      uses: styfle/cancel-workflow-action@0.9.1
      with:
        access_token: ${{ github.token }}
      if: ${{github.ref != 'refs/head/main'}}
    - uses: actions/checkout@v3
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v3
      with:
        python-version: ${{ matrix.python-version }}
    - name: Get pip cache dir
      id: pip-cache
      run: |
        python -m pip install --upgrade pip wheel
        echo "::set-output name=dir::$(pip cache dir)"
    - name: pip cache
      uses: actions/cache@v3
      with:
        path: ${{ steps.pip-cache.outputs.dir }}
        key: ${{ runner.os }}-pip-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt') }}
        restore-keys: |
          ${{ runner.os }}-pip-
    - name: Install dependencies
      run: |
        pip install -r docs/requirements.txt
    - name: Test documentation
      env:
        XLA_FLAGS: "--xla_force_host_platform_device_count=8"
      run: |
        pytest -n 1 --tb=short docs
        pytest -n 1 --tb=short --doctest-modules jax --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/interpreters/mlir.py --ignore=jax/_src/iree.py --ignore=jax/experimental/gda_serialization --ignore=jax/collect_profile.py