Stack

GitHub Actions runners for JAX

JAX wheels are GPU-specific — TPU/CUDA mismatch on default runners breaks `pjit` tests. Cirun runs your JAX CI on a right-sized cloud VM on your own account — pick AWS, GCP, Azure, Oracle, DigitalOcean, OpenStack or on-prem, billed by your cloud, never per CI minute.

Run JAX on Cirun Docs

Why this fits

JAX wheels are GPU-specific — TPU/CUDA mismatch on default runners breaks `pjit` tests.
Cirun is cloud-neutral — every supported cloud has at least one SKU that fits JAX; pick whichever account you already have.
Ephemeral by default — no state from the previous PR leaks into yours, no flaky-cache surprises.

.cirun.yml

1runners:
2  - name: jax-runner
3    cloud: aws
4    instance_type: p4d.24xlarge
5    # Use AWS Deep Learning AMI GPU PyTorch on Ubuntu 22.04, or the
6    # Cirun-published NVIDIA AMI.
7    machine_image: ami-04823729c75214919
8    labels:
9      - cirun-jax

.cirun.yml

1runners:
2  - name: jax-runner
3    cloud: gcp
4    instance_type: a2-highgpu-1g
5    # Use a GCP Deep Learning VM image family
6    # (deeplearning-platform-release) for pre-installed CUDA + drivers.
7    machine_image: projects/deeplearning-platform-release/global/images/family/pytorch-latest-cu121
8    labels:
9      - cirun-jax

.cirun.yml

1runners:
2  - name: jax-runner
3    cloud: azure
4    instance_type: Standard_ND96asr_v4
5    # Use Microsoft's HPC/DSVM Ubuntu image for pre-installed CUDA +
6    # drivers.
7    machine_image: microsoft-dsvm:ubuntu-hpc:2204:latest
8    labels:
9      - cirun-jax