jax

876b1171 - [Mosaic GPU] Support lowering SMEM slices in a warp context.

[Mosaic GPU] Support lowering SMEM slices in a warp context. Due to `warp_map` being closed from above, allocations can end up being pulled in when constructing the op. PiperOrigin-RevId: 914723326