[mosaic_gpu] Fix signedness handling in FragmentedArray._pointwise.

Only propagate signedness from operands when the output type of `op` is an `ir.IntegerType`.

PiperOrigin-RevId: 698324596

Author: petebu

jax/experimental/mosaic/gpu/fragmented_array.py

@@ -623,10 +623,6 @@ def to_layout(self, new_layout: FragmentedLayout):
     )
 
   def _pointwise(self, op, *other, output_is_signed: bool | None = None):
-    is_signed = (
-        output_is_signed if output_is_signed is not None else self.is_signed
-    )
-
     other_arrs = []
     for o in other:
       if not isinstance(o, FragmentedArray):
@@ -636,7 +632,7 @@ def _pointwise(self, op, *other, output_is_signed: bool | None = None):
           raise NotImplementedError(o)
 
         o = FragmentedArray.splat(
-            o, shape=self.shape, layout=self.layout, is_signed=is_signed
+            o, shape=self.shape, layout=self.layout, is_signed=self.is_signed
         )
 
       if isinstance(o.layout, WGSplatFragLayout):
@@ -646,7 +642,7 @@ def _pointwise(self, op, *other, output_is_signed: bool | None = None):
             o.registers.flat[0],
             shape=self.shape,
             layout=self.layout,
-            is_signed=is_signed,
+            is_signed=self.is_signed,
         )
       else:
         if self.layout != o.layout:
@@ -659,8 +655,13 @@ def _pointwise(self, op, *other, output_is_signed: bool | None = None):
 
     for idx, reg in np.ndenumerate(self.registers):
       new_regs[idx] = op(reg, *(o.registers[idx] for o in other_arrs))
+    reg_ty = new_regs.flat[0].type
+    if ir.VectorType.isinstance(reg_ty):
+      reg_ty = ir.VectorType(reg_ty).element_type
+    if output_is_signed is None and ir.IntegerType.isinstance(reg_ty):
+      output_is_signed = self.is_signed
     return FragmentedArray(
-        _registers=new_regs, _layout=self.layout, _is_signed=is_signed
+        _registers=new_regs, _layout=self.layout, _is_signed=output_is_signed
     )
 
   def __pos__(self):