Rewrite mosaic concat to support operand shapes that do not align with native shapes, Expand tests to cover multi operand, batch dim concat, etc. #62522

Workflow file for this run

.github/workflows/ci-build.yaml at f67281e

	name: CI

	# We test all supported Python versions as follows:
	# - 3.10 : Documentation build
	# - 3.10 : Part of Matrix with NumPy dispatch
	# - 3.10 : Part of Matrix
	# - 3.11 : Part of Matrix

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

	permissions:
	contents: read # to fetch code
	actions: write # to cancel previous workflows

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

	jobs:
	lint_and_typecheck:
	runs-on: ubuntu-latest
	timeout-minutes: 5
	steps:
	- uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
	- name: Set up Python 3.11
	uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
	with:
	python-version: 3.11
	- uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # ratchet: pre-commit/action@v3.0.1

	build:
	name: "build ${{ matrix.name-prefix }} (py ${{ matrix.python-version }} on ${{ matrix.os }}, x64=${{ matrix.enable-x64}})"
	runs-on: ${{ matrix.os }}
	timeout-minutes: 60
	strategy:
	matrix:
	# Test the oldest and newest supported Python versions here.
	include:
	- name-prefix: "with 3.10"
	python-version: "3.10"
	os: ubuntu-20.04-16core
	enable-x64: 1
	prng-upgrade: 1
	num_generated_cases: 1
	- name-prefix: "with 3.12"
	python-version: "3.12"
	os: ubuntu-20.04-16core
	enable-x64: 0
	prng-upgrade: 0
	num_generated_cases: 1
	steps:
	- uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
	with:
	python-version: ${{ matrix.python-version }}
	- name: Get pip cache dir
	id: pip-cache
	run: \|
	python -m pip install --upgrade pip wheel
	echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
	- name: pip cache
	uses: actions/cache@0c45773b623bea8c8e75f6c82b208c3cf94ea4f9 # ratchet: actions/cache@v4
	with:
	path: ${{ steps.pip-cache.outputs.dir }}
	key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('/setup.py', '/requirements.txt', '**/test-requirements.txt') }}
	- name: Install dependencies
	run: \|
	pip install .[minimum-jaxlib] -r build/test-requirements.txt

	- name: Run tests
	env:
	JAX_NUM_GENERATED_CASES: ${{ matrix.num_generated_cases }}
	JAX_ENABLE_X64: ${{ matrix.enable-x64 }}
	JAX_ENABLE_CUSTOM_PRNG: ${{ matrix.prng-upgrade }}
	JAX_THREEFRY_PARTITIONABLE: ${{ matrix.prng-upgrade }}
	JAX_ENABLE_CHECKS: true
	JAX_SKIP_SLOW_TESTS: true
	PY_COLORS: 1
	run: \|
	pip install -e .
	echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
	echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
	echo "JAX_ENABLE_CUSTOM_PRNG=$JAX_ENABLE_CUSTOM_PRNG"
	echo "JAX_THREEFRY_PARTITIONABLE=$JAX_THREEFRY_PARTITIONABLE"
	echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
	echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
	pytest -n auto --tb=short --maxfail=20 tests examples


	documentation:
	name: Documentation - test code snippets
	runs-on: ubuntu-latest
	timeout-minutes: 10
	strategy:
	matrix:
	python-version: ['3.10']
	steps:
	- uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
	with:
	python-version: ${{ matrix.python-version }}
	- name: Get pip cache dir
	id: pip-cache
	run: \|
	python -m pip install --upgrade pip wheel
	echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
	- name: pip cache
	uses: actions/cache@0c45773b623bea8c8e75f6c82b208c3cf94ea4f9 # ratchet: actions/cache@v4
	with:
	path: ${{ steps.pip-cache.outputs.dir }}
	key: ${{ runner.os }}-pip-docs-${{ hashFiles('/setup.py', '/requirements.txt', '**/test-requirements.txt') }}
	- name: Install dependencies
	run: \|
	pip install -r docs/requirements.txt
	- name: Test documentation
	env:
	XLA_FLAGS: "--xla_force_host_platform_device_count=8"
	JAX_TRACEBACK_FILTERING: "off"
	JAX_ARRAY: 1
	PY_COLORS: 1
	run: \|
	pytest -n auto --tb=short --doctest-glob='.md' --doctest-glob='.rst' docs --doctest-continue-on-failure --ignore=docs/multi_process.md --ignore=docs/jax.experimental.array_api.rst
	pytest -n auto --tb=short --doctest-modules jax --ignore=jax/config.py --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/_src/lib/triton.py --ignore=jax/_src/lib/mosaic_gpu.py --ignore=jax/interpreters/mlir.py --ignore=jax/experimental/array_serialization --ignore=jax/collect_profile.py --ignore=jax/_src/tpu_custom_call.py --ignore=jax/experimental/mosaic --ignore=jax/experimental/pallas --ignore=jax/_src/pallas --ignore=jax/experimental/array_api --ignore=jax/lib/xla_extension.py


	documentation_render:
	name: Documentation - render documentation
	runs-on: ubuntu-latest
	timeout-minutes: 10
	strategy:
	matrix:
	python-version: ['3.10']
	steps:
	- uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
	with:
	python-version: ${{ matrix.python-version }}
	- name: Get pip cache dir
	id: pip-cache
	run: \|
	python -m pip install --upgrade pip wheel
	echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
	- name: pip cache
	uses: actions/cache@0c45773b623bea8c8e75f6c82b208c3cf94ea4f9 # ratchet: actions/cache@v4
	with:
	path: ${{ steps.pip-cache.outputs.dir }}
	key: ${{ runner.os }}-pip-docs-${{ hashFiles('/setup.py', '/requirements.txt', '**/test-requirements.txt') }}
	- name: Install dependencies
	run: \|
	pip install -r docs/requirements.txt
	- name: Render documentation
	run: \|
	sphinx-build --color -W --keep-going -b html -D nb_execution_mode=off docs docs/build/html


	jax2tf_test:
	name: "jax2tf_test (py ${{ matrix.python-version }} on ${{ matrix.os }}, x64=${{ matrix.enable-x64}})"
	runs-on: ${{ matrix.os }}
	timeout-minutes: 30
	strategy:
	matrix:
	# Test the oldest supported Python version here.
	include:
	- python-version: "3.10"
	os: ubuntu-latest
	enable-x64: 0
	num_generated_cases: 10
	steps:
	- uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
	with:
	python-version: ${{ matrix.python-version }}
	- name: Get pip cache dir
	id: pip-cache
	run: \|
	python -m pip install --upgrade pip wheel
	echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
	- name: pip cache
	uses: actions/cache@0c45773b623bea8c8e75f6c82b208c3cf94ea4f9 # ratchet: actions/cache@v4
	with:
	path: ${{ steps.pip-cache.outputs.dir }}
	key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('/setup.py', '/requirements.txt', '**/test-requirements.txt') }}
	- name: Install dependencies
	run: \|
	pip install .[minimum-jaxlib] tensorflow -r build/test-requirements.txt

	- name: Run tests
	env:
	JAX_NUM_GENERATED_CASES: ${{ matrix.num_generated_cases }}
	JAX_ENABLE_X64: ${{ matrix.enable-x64 }}
	JAX_ENABLE_CHECKS: true
	JAX_SKIP_SLOW_TESTS: true
	PY_COLORS: 1
	run: \|
	pip install -e .
	echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
	echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
	echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
	echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
	pytest -n auto --tb=short --maxfail=20 jax/experimental/jax2tf/tests/jax2tf_test.py

	ffi:
	name: FFI example
	runs-on: ubuntu-latest
	timeout-minutes: 5
	steps:
	- uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet:actions/checkout@v4
	- name: Set up Python 3.11
	uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
	with:
	python-version: 3.11
	- name: Get pip cache dir
	id: pip-cache
	run: \|
	python -m pip install --upgrade pip wheel
	echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT
	- name: pip cache
	uses: actions/cache@0c45773b623bea8c8e75f6c82b208c3cf94ea4f9 # ratchet: actions/cache@v4
	with:
	path: ${{ steps.pip-cache.outputs.dir }}
	key: ${{ runner.os }}-pip-ffi-examples-${{ hashFiles('/setup.py', '/requirements.txt', '/test-requirements.txt', 'examples//pyproject.toml') }}
	- name: Install JAX
	run: pip install .
	- name: Build and install example project
	run: python -m pip install -v ./examples/ffi[test]
	env:
	# We test building using GCC instead of clang. All other JAX builds use
	# clang, but it is useful to make sure that FFI users can compile using
	# a different toolchain. GCC is the default compiler on the
	# 'ubuntu-latest' runner, but we still set this explicitly just to be
	# clear.
	CMAKE_ARGS: -DCMAKE_C_COMPILER=gcc -DCMAKE_CXX_COMPILER=g++
	- name: Run tests
	run: python -m pytest examples/ffi/tests

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Rewrite mosaic concat to support operand shapes that do not align with native shapes, Expand tests to cover multi operand, batch dim concat, etc. #62522

Workflow file

Rewrite mosaic concat to support operand shapes that do not align with native shapes, Expand tests to cover multi operand, batch dim concat, etc. #62522

Jobs

Run details

Workflow file for this run