jax

af36ae2c - [Pallas/Mosaic GPU] Make Mosaic GPU the default lowering backend for Pallas on GPU.

[Pallas/Mosaic GPU] Make Mosaic GPU the default lowering backend for Pallas on GPU. In order to keep using Triton for calls to `pallas_call`, setting the `backend` argument to `'triton'` is now necessary. PiperOrigin-RevId: 848532122