jax

c45ec035 - [Mosaic GPU] Initialise NVSHMEM runtime library as part of `InitKernel`.

[Mosaic GPU] Initialise NVSHMEM runtime library as part of `InitKernel`. We currently initialise the NVSHMEM runtime library before we compile the kernel. However the library is not required to be initialised for compilation. This change is required to move compilation to the "instantiate" handler where a CUDA context may not be available (e.g AOT compilation or when using TFRT). PiperOrigin-RevId: 860089154