feat: Add JAX release support v0.10.0 (#6054)

[marbre]

Cherry-picks 3b10211 into the 7.14 release branch.

---

Summary

ISSUE ID: #3878 Related PR: ROCm/rockrel#65 Extend the JAX multi-arch release workflows to support both rocm-jaxlib-v0.9.1 and rocm-jaxlib-v0.10.0 release lines.

Changes

Release workflow

• Add matrix support for:
  • rocm-jaxlib-v0.9.1
  • rocm-jaxlib-v0.10.0
• Add release-line specific metadata:
  • jax_repository
  • build_mode
  • gfx_arch

Build workflow

rocm-jaxlib-v0.9.1

• Continue using the existing native build flow.
• Build from ROCm/rocm-jax.
• Continue using build/ci_build.

rocm-jaxlib-v0.10.0

• Build from ROCm/jax.
• Build a manylinux image using the release branch Dockerfile from ROCm/rocm-jax.
• Build wheels inside the manylinux container using build/build.py.

Testing

On rockrel:

https://github.com/ROCm/rockrel/actions/runs/28217952215/job/83613621212 https://github.com/ROCm/rockrel/actions/runs/28224689859

• Preserve existing GPU-family based runner selection through configure_target_run.py.
• Add a JAX device visibility sanity check before running tests:

python -c "import jax; print(jax.local_devices())"

- Verified matrix expansion across Python 3.11–3.14.
- Verified native build path for rocm-jaxlib-v0.9.1.
- Verified manylinux build path for rocm-jaxlib-v0.10.0.
- Verified release artifact publication.
- Verified JAX device visibility check executes before tests.

TODO: 
Run JAX builds inside a top-level GitHub Actions job container, similar to existing PyTorch workflows.

[ScottTodd]

Approving for devops.

We had a successful nightly release of jax 0.10.0 that just missed the release branch cutoff: https://github.com/ROCm/rockrel/actions/runs/28346375566/job/83976592962

Note that we had a failed nightly release later at https://github.com/ROCm/rockrel/actions/runs/28355352077 (see #6210) that was from the asan release workflow and that overwrote the working packages. We reverted the nightly asan workflow in ROCm/rockrel#69 and are tracking excluding python packages and frameworks from asan releases in #6218. Future nightly releases should not have the same "asan overwrites non-asan" issue.

As for documentation, we have an updated support table at https://github.com/ROCm/TheRock/tree/main/external-builds/jax#supported-jax-versions but we are currently missing the JAX equivalent to pytorch instructions in RELEASES.md, see https://github.com/ROCm/TheRock/blob/main/RELEASES.md#installing-multi-arch-pytorch-python-packages . @erman-gurses can you or someone else update the docs?

[erman-gurses]

Approving for devops.

We had a successful nightly release of jax 0.10.0 that just missed the release branch cutoff: https://github.com/ROCm/rockrel/actions/runs/28346375566/job/83976592962

Note that we had a failed nightly release later at https://github.com/ROCm/rockrel/actions/runs/28355352077 (see #6210) that was from the asan release workflow and that overwrote the working packages. We reverted the nightly asan workflow in ROCm/rockrel#69 and are tracking excluding python packages and frameworks from asan releases in #6218. Future nightly releases should not have the same "asan overwrites non-asan" issue.

As for documentation, we have an updated support table at https://github.com/ROCm/TheRock/tree/main/external-builds/jax#supported-jax-versions but we are currently missing the JAX equivalent to pytorch instructions in RELEASES.md, see https://github.com/ROCm/TheRock/blob/main/RELEASES.md#installing-multi-arch-pytorch-python-packages . @erman-gurses can you or someone else update the docs?

Update for the RELEASES.md #6225