Fix FP8 compilation failure in jaxlib stemming from the CUDA/ROCM merge.
Fix FP8 compilation failure in jaxlib stemming from the CUDA/ROCM merge.

PiperOrigin-RevId: 484026031

google/jax,google/jax,google/jax,google/jax