Rocm jaxlib v0.9.0

[Ruturaj4]

### Problem

Pallas tests fail on AMD ROCm GPUs with the error:

ValueError: invalid literal for int() with base 10: 'gfx950'

This occurs because:

1. The Mosaic GPU backend (NVIDIA-specific) attempts to parse the GPU architecture string
2. NVIDIA GPUs return compute_capability as "major.minor" (e.g., "9.0")
3. AMD ROCm GPUs return architecture identifiers like "gfx950"
4. The code tries to parse "gfx950" as an integer, causing the failure

Solution

This PR fixes the issue with three changes:

1. jax/_src/pallas/pallas_call.py: Route ROCm devices to Triton backend

   • Added is_rocm check in gpu_lowering()
   • When backend=None and running on ROCm, automatically use Triton instead of Mosaic GPU
   • Added clear error if user explicitly requests backend='mosaic_gpu' on ROCm

2. jax/experimental/mosaic/gpu/core.py: Add safety check in _infer_arch()

   • Detect ROCm architecture strings (starting with "gfx") and raise a descriptive error

3. tests/pallas/ops_test.py: Fix test_delay skip logic

   • The delay primitive is only implemented in Mosaic GPU, not Triton
   • Updated skip condition to include ROCm devices: jtu.is_device_rocm() or not use_mosaic_gpu

Testing

• tests now pass (previously failing with gfx950 error)
• test_delay properly skips on ROCm since delay is MGPU-only

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.