[moe] Add Muon throughput experiment for grug MoE #1058

	name: Haliax - Tests

	on:
	push:
	branches:
	- main
	pull_request:

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

	jobs:
	changes:
	runs-on: ubuntu-latest
	permissions:
	contents: read
	pull-requests: read
	outputs:
	should_run: ${{ steps.filter.outputs.relevant }}
	steps:
	- uses: actions/checkout@v4
	- uses: dorny/paths-filter@v3
	id: filter
	with:
	filters: \|
	relevant:
	- 'lib/haliax/**'
	- 'uv.lock'
	- '.github/workflows/haliax-*.yaml'

	haliax-tests:
	needs: changes
	if: needs.changes.outputs.should_run == 'true'

	runs-on: ubuntu-latest

	defaults:
	run:
	working-directory: lib/haliax

	steps:
	- uses: actions/checkout@v3
	- name: Set up Python 3.11
	uses: actions/setup-python@v4
	with:
	python-version: "3.11"
	- name: Install dependencies
	run: \|
	python -m pip install uv
	uv sync --package haliax --dev
	- name: Test with pytest (JAX 0.8.0)
	run: \|
	# Test with specific JAX version
	JAX_NUM_CPU_DEVICES=8 PYTHONPATH=tests:src:. uv run --package haliax --with "jax[cpu]==0.8.0" pytest -c pyproject.toml tests

Provide feedback