Fix GitHub Actions CI: unblock auto-cancelled runs, improves test gating, adds checkpointing fixes for older JAX

[copybara-service]

Fix GitHub Actions CI: unblock auto-cancelled runs, improves test gating, adds checkpointing fixes for older JAX

Summary

The GitHub actions test suite passes on internal infrastructure but has never gone green on GitHub. This PR fixes two root causes:

1. A poisoned Actions concurrency group was auto-cancelling all runs on main since mid-June.
2. CPU-only GitHub runners failed on tests requiring TPU capabilities (int4, Pallas, scaled_matmul lowering, multi-device shardings, and TPU-specific fp8 numerics).

After these changes, the CI-equivalent suite passes locally on CPU (except two einsum tests blocked by an XLA:CPU bug). No TPU-side behavior or tolerance is changed.

1. CI Workflows

Re-keyed concurrency groups to prevent pushes to main from being starved or cancelled:

• Pushes to main: Unique group per run (github.run_id) — never queued or cancelled.
• PR runs: Shared per-PR group — superseded commits are cancelled.

concurrency:
  group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && format('pr-{0}', github.event.pull_request.number) || github.run_id }}
  cancel-in-progress: ${{ github.event_name == 'pull_request' }}

2. Test Gating

We apply targeted gating for TPU only-tests

Test File	Issue on CPU	Fix / Resolution
core/mxfp_dot_test.py	scaled_matmul MLIR lowering is missing in released PyPI JAX.	Gates the 4 affected tests using a dynamic runtime capability probe (_scaled_matmul_supported()). Tests auto-enable once JAX CPU gains lowering.
kernels/fused_hadamard_quantize_test.py	TPU-only Pallas kernel requires CpuDevice unsupported attributes during setup.	Added class-level setUp skip on non-TPU platforms.
core/conv_general_qt_test.py	Dilated fp8-backward conv (fp8_bwd_nhwc_dilated) diverges deterministically on XLA:CPU.	Skip only this specific parameterization when off-TPU.
core/dot_general_qt_test.py	Untiled fp8_bwd grads slightly exceed the strict 0.04 TPU-tuned bound on CPU.	Use a platform-conditional bound of 0.25 off-TPU (TPU keeps 0.04).

3. Checkpoint Sharding (utils/checkpoint_util.py)

Introduced a library-level fix to resolve sharding loss on released JAX versions.

Component	Details
Symptom	2 tests building NamedSharding over Mesh('fsdp': 2, 'tp': 2) with 4 virtual CPU devices passed quantized-param assertions, but the bias landed as SingleDeviceSharding.
Root Cause	On released JAX, nnx.eval_shape yields template leaves whose unboxed ShapeDtypeStruct has sharding=None. Because of this, _apply_sharding_and_dtype skipped calling jax.device_put, leaving fp params on a single device (at JAX-at-head they carry NamedSharding over an AbstractMesh).
Resolution	Implemented a metadata fallback. When the primary template-sharding resolution yields None (and use_checkpoint_sharding=False), _apply_sharding_and_dtype reconstructs the sharding from the boxed template's nnx metadata via Flax's nnx.spmd.get_var_pspec against the active concrete mesh.
Verification	Verified on JAX nightly where the primary path resolves every sharding and the helper is never invoked. Passes on both JAX 0.9.0.1 and JAX nightly (0.11.0.dev20260618).

4. Type Checks: qwix/contrib/kernels/

• Fixed pyrefly check errors on pl.kernel(...) calls.
  • Kernels are written against the newer JAX-at-head API (using out_type, scratch_types, mesh), which differs from released JAX signatures on CPU runners.