jax

1efef6bf - [pallas:mosaic_gpu] `emit_pipeline` now correctly supports `BlockSpec`s in GMEM

[pallas:mosaic_gpu] `emit_pipeline` now correctly supports `BlockSpec`s in GMEM This is necessary to replace the pipelining logic in the lowering with `emit_pipeline`. PiperOrigin-RevId: 698858380