jax

59ae2af6 - [pallas:mosaic_gpu] Added a test doing manual in kernel pipelining

[pallas:mosaic_gpu] Added a test doing manual in kernel pipelining I think we have most of the primitives necessary, so the next step is to sketch `emit_pipeline`. PiperOrigin-RevId: 684840800