jax-ml/jax · error · ValueError

Incompatible FragmentedArray shapes

Error message

Incompatible FragmentedArray shapes

What it means

Even with identical layouts, two FragmentedArrays can hold different numbers of register tiles (registers_shape depends on shape and layout parameters). _pointwise pairs registers index-by-index, so unequal register arrays cannot be combined and raise ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:1507

        o = FragmentedArray.splat(
            o, shape=self.shape, layout=self.layout, is_signed=self.is_signed
        )

      if isinstance(o.layout, WGSplatFragLayout):
        if not o.layout.can_broadcast_to(self.shape):
          raise ValueError(
              f"Cannot broadcast shape {self.shape} to layout {o.layout}")
        o = FragmentedArray.splat(
            o.registers.flat[0],
            shape=self.shape,
            layout=self.layout,
            is_signed=o.is_signed,
        )
      else:
        if self.layout != o.layout:
          raise ValueError("Incompatible FragmentedArray layouts")
        if self.registers.shape != o.registers.shape:
          raise ValueError("Incompatible FragmentedArray shapes")

      other_arrs.append(o)
    new_regs = np.empty_like(self.registers)

    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 isinstance(reg_ty, ir.VectorType):
      reg_ty = ir.VectorType(reg_ty).element_type
    if output_is_signed is None and isinstance(reg_ty, ir.IntegerType):
      output_is_signed = self.is_signed
    return FragmentedArray(
        _registers=new_regs, _layout=self.layout, _is_signed=output_is_signed
    )

  def __pos__(self):
    return self

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/pad both operands to the identical shape before the pointwise op
2. If broadcasting was intended, splat the smaller operand to the larger shape (FragmentedArray.splat / broadcast) instead of relying on NumPy rules
3. Assert o.registers.shape == self.registers.shape in tests to catch silent tile-shape divergence

Example fix

# before
z = a + b  # a.shape=(128,8), b.shape=(128,16)
# after
b = b[:, :8]  # or pad a to (128,16)
z = a + b

assert self.registers.shape == o.registers.shape, (
    f"register shapes differ: {self.registers.shape} vs {o.registers.shape}")

def same_registers(a, b) -> bool:
    return a.registers.shape == b.registers.shape

try:
    z = a + b
except ValueError as e:
    if 'Incompatible FragmentedArray shapes' in str(e):
        raise ValueError(f"pad/reshape operands {a.shape} and {b.shape} to match") from e
    raise

When it happens

Trigger: Applying arithmetic between arrays with the same layout but different element shapes — e.g. (128, 8) + (128, 16) in the same layout — so registers.shape differs; also shapes that are broadcast-compatible in NumPy but not in register-space Mosaic.

Common situations: Assuming NumPy broadcasting semantics for Mosaic fragments; padding one operand's tile shape differently; partial-tile kernels where one operand was padded to a full tile and the other was not.

Related errors

• Swizzle {self.swizzle} requires the trailing dimension to be
• grid_names must have the same length as grid, got {self}.
• cluster_names must have the same length as cluster, got {sel
• {tiling=} and {grid=} must have same length.
• The stored value has shape {src.shape}, but the target refer

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/c5656f4cb156eb78. Report an issue: GitHub.