(1/N) Add gmm v2 to Tokamax ragged dot API. #2657

	name: CI-experimental

	on:
	# Trigger the workflow on push or pull request,
	# but only for the main branch
	push:
	branches:
	- main
	pull_request:
	branches:
	- main
	# Event manually triggered (in the Github UI)
	workflow_dispatch:

	# Constants for the workflow that are shared across jobs.
	env:
	JAX_ENABLE_X64: 0
	JAX_TRACEBACK_FILTERING: off

	permissions:
	contents: read
	# We should cancel the previous run if it's still running and there is a new commit.
	concurrency:
	group: ${{ github.workflow }}-${{ github.head_ref \|\| github.ref }}
	cancel-in-progress: true

	# These jobs should build Tokamax and run the tests on Google Cloud resources.
	jobs:
	tokamax-experimental-tests:
	strategy:
	matrix:
	runner: ['linux-x86-ct6e-44-1tpu']
	runs-on: ${{ matrix.runner }}
	defaults:
	run:
	shell: bash
	container:
	# Docker image will include CUDA and cuDNN for Tokamax.
	image: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda13.2-cudnn9.15'
	timeout-minutes: 300
	steps:
	- name: Checkout Tokamax
	uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
	with:
	persist-credentials: false

	- name: Install dependencies
	run: \|
	python3.12 -m uv pip install --upgrade pip
	python3.12 -m uv pip --version
	python3.12 -m uv pip install -e .[tpu,test]
	python3.12 -m uv pip freeze
	- name: Test Tokamax ops with unittest
	run: \|
	pytest -s --ignore-glob="/test_base.py" --ignore-glob="/splash_attention_kernel_sharded_test.py" tokamax/_src/ops/experimental/tpu

Provide feedback