Skip to content

Feature/batched vmap #2214

Feature/batched vmap

Feature/batched vmap #2214

Workflow file for this run

name: Tests
on:
schedule:
- cron: 00 00 * * 1
push:
branches: [main]
pull_request:
branches: [main]
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
fast-tests:
name: Fast tests Python ${{ matrix.python-version }} ${{ matrix.jax-version }}
runs-on: ubuntu-latest
# allow tests using the latest JAX to fail
continue-on-error: ${{ matrix.jax-version == 'jax-latest' }}
strategy:
fail-fast: false
matrix:
python-version: ['3.10']
jax-version: [jax-default, jax-latest]
steps:
- uses: actions/checkout@v4
- name: Set up Python 3.9
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install tox
- name: Setup environment
run: |
tox -e py${{ matrix.python-version }}-${{ matrix.jax-version }} --notest -v
- name: Run tests
run: |
tox -e py${{ matrix.python-version }}-${{ matrix.jax-version }} --skip-pkg-install -- -m fast --memray -n auto -vv
gpu-tests:
name: Fast GPU tests Python 3.10 on ubuntu-22.04
runs-on: [self-hosted, ott-gpu]
container:
image: docker://michalk8/cuda:12.2.2-cudnn8-devel-ubuntu22.04
options: --gpus="device=2"
steps:
- uses: actions/checkout@v4
- name: Install dependencies
run: |
python3 -m pip install --upgrade pip
python3 -m pip install -e".[test]"
python3 -m pip install "jax[cuda12]"
- name: Run nvidia-smi
run: |
nvidia-smi
- name: Run tests
run: |
python3 -m pytest -m "fast and not cpu" --memray --durations 10 -vv
env:
XLA_PYTHON_CLIENT_PREALLOCATE: 'false'
tests:
name: Python ${{ matrix.python-version }} on ${{ matrix.os }}
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
python-version: ['3.10', '3.11', '3.12', '3.13']
os: [ubuntu-latest]
include:
- python-version: '3.9'
os: macos-13
- python-version: '3.10'
os: macos-14
steps:
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install tox
- name: Setup environment
run: |
tox -e py${{ matrix.python-version }} --notest -v
- name: Run tests
run: |
tox -e py${{ matrix.python-version }} --skip-pkg-install
env:
PYTEST_ADDOPTS: --memray -vv
- name: Upload coverage
uses: codecov/codecov-action@v4
with:
files: ./coverage.xml
flags: tests-${{ matrix.os }}-${{ matrix.python-version }}
name: unittests
token: ${{ secrets.CODECOV_TOKEN }}
env_vars: OS,PYTHON
fail_ci_if_error: false
verbose: true