Skip to content

Collapse jax build/test double-counting by normalizing jax_ref to bare version - #84

Open
HereThereBeDragons wants to merge 2 commits into
developfrom
users/lpromber/normalize_jax_ref
Open

Collapse jax build/test double-counting by normalizing jax_ref to bare version#84
HereThereBeDragons wants to merge 2 commits into
developfrom
users/lpromber/normalize_jax_ref

Conversation

@HereThereBeDragons

Copy link
Copy Markdown
Collaborator

Summary

We fix a double-counting bug in the release status.json where every jax build
and test cell is counted twice. The cause is that one jax version reaches the
status pipeline under two spellings: the build matrix names cells by git ref
(rocm-jaxlib-v0.11.0), while the release orchestrator and test dispatch inputs
use the bare version (0.11.0). Left un-normalized, the two spellings key two
distinct variants that never merge, doubling the jax counts.

We normalize every jax ref to the bare version at the two points a ref enters a
variant key, so the two spellings collapse to a single (py, version) cell. On
the affected release this reduces the jax variant count from 30 to 15.

Change

  • Add _normalize_ref(axis_key, ref), called uniformly for every fan-out axis.
    Today only jax needs it: a jax_ref value has the rocm-jaxlib-v prefix
    stripped, and the torch axis is a pure passthrough.
  • Wire it into both variant-key entry points: _variants_from_jobs
    (job-name parsing) and _variants_from_inputs (dispatch-input parsing).

We strip the prefix rather than add it because stripping is idempotent on
already-bare refs and never mangles a non-version ref such as a branch name.

Tests

  • A jax cell derived from a prefixed build job and a bare-version test input now
    collapses to one variant with a single key.
  • A torch ref (including slashes) passes through normalization verbatim,
    confirming the torch axis is unaffected.

Scope

The fix applies to future releases only. The already-frozen
20260820/status.json is not retroactively corrected.

@HereThereBeDragons
HereThereBeDragons requested review from a team and marbre August 21, 2026 14:41
@HereThereBeDragons
HereThereBeDragons changed the base branch from main to develop August 21, 2026 14:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant