jax

c965ffbf - [Mosaic GPU] Remove expect_wait from Barrier.wait

[Mosaic GPU] Remove expect_wait from Barrier.wait It looks like LLVM already moves the wait loops to the end of the program, so the whole optimization is no longer necessary and only adds unnecessary operations. PiperOrigin-RevId: 703052393