jax

745cbe2f - [Mosaic GPU][NFC] Pass `CompiledKernel` by pointer to `CachedInit`.

[Mosaic GPU][NFC] Pass `CompiledKernel` by pointer to `CachedInit`. We change the kernel type from const reference to const pointer to make it clear that `CachedInit` requires pointer stability. PiperOrigin-RevId: 861723228