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

	name: JAX Array API

	on:
	push:
	branches:
	- main
	pull_request:
	branches:
	- main

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

	jobs:
	build:

	runs-on: ubuntu-latest
	strategy:
	matrix:
	python-version: [3.11]

	steps:
	- name: Checkout jax
	uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet: actions/checkout@v4
	- name: Checkout array-api-tests
	uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # ratchet: actions/checkout@v4
	with:
	repository: data-apis/array-api-tests
	# TODO(jakevdp) update this to a stable release/tag when available.
	ref: 'b4c0823469c02d6ce6e512ad4c2bd8ba42b1b4b2' # Latest commit as of 2024-09-09
	submodules: 'true'
	path: 'array-api-tests'
	- name: Set up Python ${{ matrix.python-version }}
	uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3 # ratchet:actions/setup-python@v5
	with:
	python-version: ${{ matrix.python-version }}
	- name: Install dependencies
	run: \|
	python -m pip install .[ci]
	python -m pip install -r array-api-tests/requirements.txt
	- name: Run the test suite
	env:
	ARRAY_API_TESTS_MODULE: jax.numpy
	JAX_ENABLE_X64: 'true'
	run: \|
	cd ${GITHUB_WORKSPACE}/array-api-tests
	pytest array_api_tests --max-examples=5 --derandomize --disable-deadline --skips-file ${GITHUB_WORKSPACE}/tests/array_api_skips.txt

Provide feedback