jax

51d35912 - [Mosaic GPU] Reenable the use of redux.sync, but only when all lanes participate

[Mosaic GPU] Reenable the use of redux.sync, but only when all lanes participate This seems to be the only case when this instruction is actually profitable. PiperOrigin-RevId: 857092650