jax

6eaada72 - [Mosaic GPU] Run the kernel in `CustomCallTest.MosaicGpuUsesCommandBuffers`.

[Mosaic GPU] Run the kernel in `CustomCallTest.MosaicGpuUsesCommandBuffers`. We adjust the custom calls in the test to pass in valid MGPU kernels. PiperOrigin-RevId: 859517397