[Mosaic GPU] Expose wait_parity on collective barrier

PiperOrigin-RevId: 666761011
jax-ml · Aug 23, 2024 · c767875 · c767875
1 parent c430b0c
commit c767875
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
@@ -712,8 +712,11 @@ def arrive(self):
  has_side_effects=True,
  )
 
- def wait(self):
- self.barrier.wait()
+ def wait(self, *args, **kwargs):
+ self.barrier.wait(*args, **kwargs)
+
+ def wait_parity(self, *args, **kwargs):
+ self.barrier.wait_parity(*args, **kwargs)
 
 
 class Partition: