jax_tests.yml

name: Dependency test JAX

on:
  pull_request:
    types: [labeled]
  workflow_dispatch:

jobs:
  jax_tests:
    if: ${{ github.event.label.name == 'test_jax' && github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' }}
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        jax-version: [0.3.0,  0.3.1,  0.3.2,  0.3.3,  0.3.4,  0.3.5,
                      0.3.6,  0.3.7,  0.3.8,  0.3.9,  0.3.10, 0.3.11,
                      0.3.12, 0.3.13, 0.3.14, 0.3.15, 0.3.16, 0.3.17,
                      0.3.19, 0.3.20, 0.3.21, 0.3.22, 0.3.23, 0.3.24,
                      0.3.25, 0.4.1,  0.4.2,  0.4.3,  0.4.4,  0.4.5,
                      0.4.6,  0.4.7,  0.4.8,  0.4.9,  0.4.10, 0.4.11,
                      0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18,
                      0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23,
                      0.4.24, 0.4.25]
        group: [1, 2]
    steps:
      - uses: actions/checkout@v4
      - name: Set up Python 3.9
        uses: actions/setup-python@v5
        with:
          python-version: 3.9
          cache: pip
      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          pip install -r devtools/dev-requirements.txt
          pip install matplotlib==3.5.0
      - name: Remove jax
        run: |
          pip uninstall jax jaxlib -y
      - name: install jax
        run: |
          pip install "jax[cpu]==${{ matrix.jax-version }}"
      - name: Test with pytest
        run: |
          pwd
          lscpu
          python -m pytest -m unit \
            --durations=0 \
            --mpl \
            --maxfail=1 \
            --splits 3 \
            --group ${{ matrix.group }} \
            --splitting-algorithm least_duration