jax

ce449dc0 - [Pallas/Mosaic GPU] Fix misuse of `self.{kernel,pallas_call}` in `mosaic_gpu_test.py`.

[Pallas/Mosaic GPU] Fix misuse of `self.{kernel,pallas_call}` in `mosaic_gpu_test.py`. Also monkey patch the modules locally, and make sure that the problematic functions can no longer be called by test writers without raising an `Exception`. PiperOrigin-RevId: 883079468