jax

4d786dec - [Mosaic GPU] Move GPU module loading and function resolution out of runtime.

[Mosaic GPU] Move GPU module loading and function resolution out of runtime. This change removes the required compilation of the initialization function for GPU kernels and loads modules directly in custom call handling. * Extract GPU binary and kernel metadata (name, dynamic shared memory size, and cluster size) as module attributes during launch lowering instead of generating a compiled initialization function. * Serialize the GPU binary and metadata directly in MosaicGpuKernelProto. * Invoke CUDA Driver and NVSHMEM APIs directly from custom_call.cc during kernel initialization. * Remove mosaic_gpu_module_load and mosaic_gpu_get_function from runtime.cc. We bump the mosaic_gpu proto version which will trigger XLA recompilation if there is a version mismatch. PiperOrigin-RevId: 955981262