Collapse jax build/test double-count and rename axis to jax_version

[HereThereBeDragons]

Summary

We fix a double-counting bug in the release status.json where every jax build and test cell is counted twice. 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, jax_version) cell. On the affected release this reduces the jax variant count from 30 to 15.

We also rename the jax variant axis jax_ref -> jax_version. After normalization the axis holds a bare version, not a git ref, so the old name was inaccurate. The new name matches TheRock's own explicit jax_version input, which is now passed alongside jax_ref.

Why not keep the full ref

The full rocm-jaxlib-v... spelling exists only on the build side. A test job's name carries only the bare-version ancestor, and the per-cell test inputs are not recoverable because all cells nest under one release run_id (notify inputs collapse to a single cell). Verified on run 32806240964: 14/14 build job names carry the full ref, 0/14 test job names do. Bare is therefore the only spelling common to both sides, so stripping is the only reconciliation that reads real data rather than synthesizing a prefix onto refs the test side never reports.

Change

• Add _normalize_ref(axis_key, ref), called uniformly for every fan-out axis. Today only jax needs it: a jax value has the rocm-jaxlib-v prefix stripped; 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).
• Rename the axis key to jax_version in _VARIANT_AXIS_KEY, _VARIANT_INPUT_KEYS, and _normalize_ref's dispatch.
• Prefer TheRock's explicit bare jax_version input over jax_ref when both are present; jax_ref is still read and stripped when it is the only input.
• Update docs/status-json/status_json_reference.jsonc examples to jax_version (bare).

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 collapses to one variant with a single key.
• A torch ref (including slashes) passes through normalization verbatim, confirming the torch axis is unaffected.
• A jax test variant reads and strips jax_ref when it is the only input.
• A jax test variant prefers the explicit jax_version input over jax_ref when both are supplied.

[marbre]

From my review agent:

Preserve jax_ref as a Git ref — /develop/quartz/scripts/receive_therock/therock_update_status_json.py:382

_normalize_ref() strips rocm-jaxlib-v from every JAX ref, including dispatch inputs. TheRock deliberately distinguishes the checkoutable jax_ref from the bare display-only jax_label (/develop/therock/ build_tools/github_actions/configure_jax_release_matrix.py:122), and Quartz’s schema reference publishes the full ref (/develop/quartz/docs/status-json/status_json_reference.jsonc:263).

This changes the meaning of a public schema-v2 field and can alias distinct valid refs—for example, rocm-jaxlib-vfeature and feature become the same variant key. I recommend canonicalizing the bare job label toward the full ref, or separating display normalization from variant identity, while preserving arbitrary branch/tag/SHA inputs.

[HereThereBeDragons]

thanks. that turned out to be more tricky than expected.

i renamed now the parameter from "jax_ref" to "jax_version" to clarify this.
we cannot use rocm-jaxlib-v.. as:

• jax test only carry the version but not the ref/branch in the job title
• we use the job list to update the cells (so we can also update the build jobs on the fly when the test jobs are reporting in), and not based on the workflow inputs - thus we need to use the job title

as such using the full jax_ref with rocm-jaxlib-v.. could be misleading and wrong if we extrapolate it for a test job and jax was build using a different ref.