jax

65cd8839 - [NFC] [Pallas TPU] Refactor Pallas Mosaic lowering to support unpipelined kernels via a simpler lowering path

[NFC] [Pallas TPU] Refactor Pallas Mosaic lowering to support unpipelined kernels via a simpler lowering path This change introduces separate lowering paths for pipelined and unpipelined kernels in Pallas Mosaic. The MPMD lowering now uses the unpipelined path. PiperOrigin-RevId: 905273214