jax

482ae677 - [pallas:mosaic] Added a `register_lowering` decorator

[pallas:mosaic] Added a `register_lowering` decorator This change prepares the switch to `pltpu.KernelType`-specific lowering rules. PiperOrigin-RevId: 755864420