[pallas:mosaic_gpu] FragmentedArray.reduce_sum now returns a FragmentedArray

This aligns it with the `reduce` method and also makes it clear that the
reduction always produces a scalar.

PiperOrigin-RevId: 703494443

Author: superbobry

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -1256,13 +1256,11 @@ def _reduce_sum_lowering_rule(ctx: LoweringRuleContext, x, *, axes):
   [x_aval] = ctx.avals_in
   match x.layout:
     case mgpu.WGStridedFragLayout():
-      if axes != (0,):
-        raise NotImplementedError("No support for axes other than 0 yet")
+      if set(axes) != set(range(x_aval.ndim)):
+        raise NotImplementedError("No support for axes yet")
       scratch_ty = jax.ShapeDtypeStruct(shape=(4,), dtype=x_aval.dtype)
       with ctx.module_ctx.scratch_view([scratch_ty]) as [scratch]:
-        return mgpu.FragmentedArray.splat(
-            x.reduce_sum(scratch), (), is_signed=mgpu_utils.is_signed(x_aval.dtype)
-        )
+        return x.reduce_sum(scratch)
     case mgpu.WGMMA_LAYOUT:
       if axes != (x_aval.ndim - 1,):
         raise NotImplementedError

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -1129,7 +1129,7 @@ def upcast_to_bf16(reg, high):
     )
 
   # NOTE: scratch can be reused immediately once this function returns.
-  def reduce_sum(self, scratch) -> ir.Value:
+  def reduce_sum(self, scratch):
     if ir.FloatType.isinstance(self.mlir_dtype):
       op = addf
     elif ir.IntegerType.isinstance(self.mlir_dtype):
@@ -1168,7 +1168,7 @@ def reduce_sum(self, scratch) -> ir.Value:
     utils.warpgroup_barrier()
     result = memref.load(scratch, [zero_index])
     utils.warpgroup_barrier()  # Make sure everyone is done using scratch.
-    return result
+    return FragmentedArray.splat(result, (), is_signed=self.is_signed)
 
   def reduce(self, op: str | Callable[[ir.Value, ir.Value], ir.Value], axis):
     if isinstance(op, str):

tests/mosaic/gpu_test.py

@@ -1481,11 +1481,7 @@ def kernel(ctx, src, dst, scratch):
       src = mgpu.FragmentedArray.load_strided(
           src, is_signed=utils.is_signed(dtype)
       )
-      acc = mgpu.FragmentedArray.splat(
-          src.reduce_sum(scratch),
-          (m,),
-          is_signed=src.is_signed
-      )
+      acc = src.reduce_sum(scratch).broadcast((m,))
       acc.store_untiled(dst)
 
     in_shape = jax.ShapeDtypeStruct((m, n), dtype)