Skip to content

Rename from jax_ml_stack to jax-ai-stack #40

Rename from jax_ml_stack to jax-ai-stack

Rename from jax_ml_stack to jax-ai-stack #40

Workflow file for this run

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
# TODO(jakevdp): add testing on macOS-11 and windows-2019 when grain supports them,
# or alternatively run a subset of tests without grain.
jobs:
build-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"]
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 -U jax flax grain ml_dtypes optax orbax tensorflow tensorflow_datasets pytest pytest-xdist
- name: Run tests
run: |
pytest -n auto jax_ai_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,tfds,grain]
- name: Run tests
run: |
pytest -n auto jax_ai_stack