Add ROCm pytest workflow to nightly/release CI
Integrate pytest_rocm.yml workflow into wheel_tests_nightly_release.yml
for nightly and release testing. The workflow tests ROCm functionality
across multiple configurations:
Matrix strategy:
- GPU configurations: 1-GPU and 4-GPU AMD systems
- Python versions: 3.11, 3.12
- ROCm version: 7.2.0 (container tag: rocm720)
- Ubuntu: Fixed to 24.04 (decoupled from Python version)
Configuration:
- jaxlib source: head (from GCS bucket)
- ROCm plugins: Latest release from ROCm/jax-plugin
- Respects skip-download-jaxlib-and-plugins-from-gcs input
- Uses configurable gcs_download_uri input
The workflow name dynamically reflects whether testing nightly or release
artifacts based on the branch name (release/* branches use "latest release",
others use "nightly").
Also updates pytest_rocm.yml default GCS URI to use static
gs://jax-nightly-artifacts/latest instead of dynamic workflow context.
