[CI] Pin JAX image to 2026-07-21

[fheinecke]

Description

CI tests on GH runners are failing because they pull the latest, rolling ghcr.io/nvidia/jax:jax image tag. This poses a couple of problems:

• You never really know what you're testing against until the image is pulled (varies day to day)
• When upstream pushes a broken release (see last night's release where only an arm64 image was pushed), it breaks CI pipelines

This locks the image to ghcr.io/nvidia/jax:jax-2026-07-21. The tradeoff here is that the image will need to be periodically updated. IMO this should be pinned further, but locking it to the 07/21 image should get CI unblocked.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Lock the jax image to ghcr.io/nvidia/jax:jax-2026-07-21

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[fheinecke]

The important job to watch here is the "All" job. The JAX job has a pre-existing failure that has surfaced again

[greptile-apps]

Greptile Summary

Pins the JAX container image used in CI from the rolling ghcr.io/nvidia/jax:jax tag to the specific jax-2026-07-21 date tag, unblocking CI after an upstream arm64-only push broke the amd64 pipeline.

• Both references in build.yml — the container.image field for the JAX job and the docker run command inside the wheel-build job — are updated consistently.
• The date-pinned tag is a temporary measure; a follow-up to either use a digest or adopt a regular update process would improve long-term reproducibility.

Confidence Score: 5/5

Safe to merge — pins a rolling image tag to a known-good dated release, unblocking CI with no logic changes.

Both image references in the workflow are updated to the same pinned tag, all other workflow files are unaffected, and the change has no impact outside CI infrastructure.

Files Needing Attention: No files require special attention.

Important Files Changed

Filename	Overview
.github/workflows/build.yml	Both JAX image references updated from the rolling jax tag to the pinned jax-2026-07-21 date tag consistently.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[CI Trigger] --> B{Job Type}
    B --> C[JAX Build Job\ncontainer: ghcr.io/nvidia/jax:jax-2026-07-21]
    B --> D[Wheel Build Job]
    D --> E[docker run\nghcr.io/nvidia/jax:jax-2026-07-21]
    C --> F[Install Dependencies & Run Tests]
    E --> G[Build Wheel Inside Container]

Loading

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into fred/pin-ci-jax..." | Re-trigger Greptile}

[timmoon10]

This makes sense to me. Ideally we'd use official releases (like the NGC container), but in practice it seems JAX development is less stable than PyTorch and it's easier to deal with breaking changes by pinning to a nightly build than waiting for the next monthly release.

[tdophung]

Hi, do you intend to update the jax image once per month?

[ksivaman]

/te-ci jax L0 L1

[ksivaman]

LGTM