test.yaml

name: Test

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

permissions:
  contents: read  # to fetch code

jobs:
  built-latest:
    name: Latest packages (${{ matrix.os }} Python ${{ matrix.python-version }})
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        os: ["ubuntu-latest"]
        python-version: ["3.10", "3.11", "3.12"]
        include:
        - os: windows-2019
          python-version: "3.11"

    steps:
    - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
      with:
        submodules: true
    - 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 --upgrade pip
        pip install .[dev]
        pip install -U jax flax optax orbax
    - name: Run tests
      run: |
        pytest -n auto jax_ml_stack

  build:
    name: Pinned packages (${{ matrix.os }} Python ${{ matrix.python-version }})
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        os: ["ubuntu-latest"]
        python-version: ["3.10", "3.11", "3.12"]

    steps:
    - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
      with:
        submodules: true
    - 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 --upgrade pip
        pip install .[dev]
    - name: Run tests
      run: |
        pytest -n auto jax_ml_stack