[Mosaic GPU] Fix broken tests when jax is coupled with an old jax lib version.

PiperOrigin-RevId: 840181806

Author: dimitar-asenov

jax/_src/pallas/mosaic_gpu/primitives.py

@@ -249,21 +249,31 @@ def _copy_smem_to_gmem_lowering(
     )
   assert not copy_params.get("gmem_transform")
   if reduction_op is not None:
+    # TODO(b/415721295): Call mgpu.dialect.async_store after the if, after
+    # the minimal jaxlib version is 0.8.2.
+    if not hasattr(mgpu.dialect, "TMAReduction"):
+      raise NotImplementedError("Reduction op is not supported yet.")
     reduction_op_attr = getattr(
         mgpu.dialect.TMAReduction, reduction_op.capitalize()
     )
+    mgpu.dialect.async_store(
+        src,
+        dst,
+        indices,
+        slice_lengths,
+        predicate=predicate,
+        commit_group=commit_group,  # type: ignore[call-arg]
+        reduction_op=reduction_op_attr,
+    )
   else:
-    reduction_op_attr = None
-
-  mgpu.dialect.async_store(
-      src,
-      dst,
-      indices,
-      slice_lengths,
-      predicate=predicate,
-      commit_group=commit_group,  # type: ignore[call-arg]
-      reduction_op=reduction_op_attr,
-  )
+    mgpu.dialect.async_store(
+        src,
+        dst,
+        indices,
+        slice_lengths,
+        predicate=predicate,
+        commit_group=commit_group,  # type: ignore[call-arg]
+    )
   return ()
 
 

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -991,7 +991,8 @@ def _mgpu_async_store_op_lowering_rule(
   # flatten -> async_copy -> unflatted here, as long as flattened size is a
   # multiple of 16.
 
-  if store_op.reduction_op is not None:
+  # TODO(b/415721295):Simplify, after the minimal jaxlib version is 0.8.2.
+  if hasattr(mgpu, "TMAReduction") and store_op.reduction_op is not None:
     reduction_op = mgpu.TMAReduction(store_op.reduction_op.value).name.lower()
   else:
     reduction_op = None

tests/mosaic/gpu_test.py

@@ -5348,6 +5348,10 @@ def body(
 
   @parameterized.parameters(jnp.float32, jnp.bfloat16, jnp.float16)
   def test_async_store_add_reduction(self, dtype):
+    # TODO(b/415721295):Remove after the minimal jaxlib version is 0.8.2.
+    if not hasattr(mgpu_dialect, "TMAReduction"):
+      self.skipTest("TMAReduction op is required.")
+
     shape = (8, 128)
 
     def body(ctx, src, dst, smem):

tests/pallas/mosaic_gpu_test.py

@@ -654,6 +654,10 @@ def kernel(x_ref, o_ref_gmem, scratch_ref):
 
   @parameterized.parameters(jnp.bfloat16, jnp.float16, jnp.float32)
   def test_copy_smem_to_gmem_reduction(self, dtype):
+    # TODO(b/415721295):Remove after the minimal jaxlib version is 0.8.2.
+    if not hasattr(mgpu.dialect, "TMAReduction"):
+      self.skip_if_wg_semantics()
+
     @functools.partial(
         self.pallas_call,
         grid=(200,),