[Mosaic_GPU] Fix a typo that caused a regression

PiperOrigin-RevId: 647347033
google · Jun 27, 2024 · 61185a2 · 61185a2
1 parent 9ae1c56
commit 61185a2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/experimental/mosaic/gpu/wgmma.py b/jax/experimental/mosaic/gpu/wgmma.py
@@ -51,7 +51,7 @@ def __init__(self, *, _value: mgpu.FragmentedArray, _sync: bool = True):
       raise ValueError("Only WGMMA layouts supported in WGMMAAccumulator")
     self.value = _value
     if _sync:
-      self._value = wgmma_fence(_value)
+      self.value = wgmma_fence(_value)
 
   @classmethod
   def zero(cls, m, n, dtype=None):