jax

30eaa17e - [Mosaic GPU] Use existing CUDA compute capability whenever available.

[Mosaic GPU] Use existing CUDA compute capability whenever available. We currently query CUDA compute capability from the GPU device. However we want to support use cases where this information can't be queried (e.g. for AOT compilation no CUDA context is available). We fallback to the old mechanism for backward compatibility with the legacy custom call API. PiperOrigin-RevId: 860059274