jax

3d87a01b - [Pallas:MGPU] Adjust warpgroup lowering to the recent emit_pipeline changes

[Pallas:MGPU] Adjust warpgroup lowering to the recent emit_pipeline changes The Pallas-level pipelining generates a number of ops we haven't had to deal with before like conditionals, scans, etc. PiperOrigin-RevId: 730899808