jax

04e4c69f - [mosaic_gpu] Handle older `jaxlib`s in the profiler module

[mosaic_gpu] Handle older `jaxlib`s in the profiler module `measure` now raises a `RuntimeError` if the available `jaxlib` does not have the required custom calls. PiperOrigin-RevId: 698351662