[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

This PR replaces the rolling ghcr.io/nvidia/jax:jax container tag with a date-pinned tag ghcr.io/nvidia/jax:jax-2026-07-21 to restore CI stability after an upstream release pushed an arm64-only image that broke GitHub Actions runners.

• Both references to the JAX image are updated consistently: the container-level image in the jax job and the docker run call in the all job.
• The change is a short-term fix acknowledged by the author; longer-term, pinning to an image digest (SHA) rather than a mutable date tag would offer stronger reproducibility guarantees.

Confidence Score: 5/5

Safe to merge — the change is a targeted two-line update that locks the JAX image to a known-good date tag and directly addresses the rolling-tag breakage in CI.

Both usages of the JAX image tag are updated consistently (the job-level container and the docker run call), and the replaced rolling tag was the confirmed cause of CI failures. No logic, test, or application code is touched.

No files require special attention. The only suggestion is a follow-up hardening step (digest pinning) noted by the PR author themselves.

Important Files Changed

Filename	Overview
.github/workflows/build.yml	Pins the JAX container image from the rolling jax tag to jax-2026-07-21 in both the jax job container and the all job's docker run command — change is minimal, consistent, and correctly addresses the rolling-tag instability.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[build.yml workflow] --> B[core job]
    A --> C[pytorch job]
    A --> D[jax job]
    A --> E[all job]

    B --> B1["nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04\n(unchanged)"]
    C --> C1["nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04\n(unchanged)"]
    D --> D1["ghcr.io/nvidia/jax:jax-2026-07-21\n(pinned from :jax rolling tag)"]
    E --> E1["ghcr.io/nvidia/jax:jax-2026-07-21\n(pinned from :jax rolling tag)"]

    style D1 fill:#90EE90
    style E1 fill:#90EE90

Loading

_{Reviews (1): Last reviewed commit: "[CI] Pin JAX image to 2026-07-21" | Re-trigger Greptile}

[greptile-apps]

Date tags are mutable — an image registry owner can push a new image to jax-2026-07-21 at any time. Pinning to an immutable digest (e.g. ghcr.io/nvidia/jax@sha256:<digest>) would guarantee bit-for-bit reproducibility. Since the PR description already notes this is a temporary measure, this is just a pointer for the follow-up.

Suggested change

	image: ghcr.io/nvidia/jax:jax-2026-07-21
	image: ghcr.io/nvidia/jax@sha256:<digest-of-jax-2026-07-21> # pin by digest for immutability

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!