Test

Update JAX AI for PyTorch users tutorial #432

	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

	concurrency:
	group: ${{ github.workflow }}-${{ github.head_ref \|\| github.ref }}
	cancel-in-progress: true

	# 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:
	name: Pinned packages (${{ matrix.os }} Python ${{ matrix.python-version }})
	runs-on: ${{ matrix.os }}
	timeout-minutes: 10
	strategy:
	matrix:
	os: ["ubuntu-latest"]
	python-version: ["3.10", "3.11", "3.12"]

	steps:
	- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
	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

Provide feedback