jax-ml/jax · error · ValueError

Incompatible FragmentedArray layouts

Error message

Incompatible FragmentedArray layouts

What it means

_pointwise requires both operands to end up in the identical fragment layout so registers can be paired elementwise. After splat handling, if self.layout != o.layout for a non-splat operand, there is no implicit conversion and the op aborts with ValueError.

Source

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

          raise NotImplementedError(o)

        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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Explicitly convert one operand: other = other.to_layout(self.layout) before the op
2. Produce both operands with the same layout at creation time (same load/layout parameters)
3. Print/assert self.layout == other.layout in kernel debug builds to catch drift early

Example fix

# before
z = acc + rhs  # acc.layout != rhs.layout
# after
z = acc + rhs.to_layout(acc.layout)

assert self.layout == o.layout, f"layout mismatch: {self.layout} vs {o.layout}"

def same_layout(a, b) -> bool:
    return a.layout == b.layout

try:
    z = a + b
except ValueError as e:
    if 'Incompatible FragmentedArray layouts' in str(e):
        z = a + b.to_layout(a.layout)
    else:
        raise

When it happens

Trigger: Applying +, -, *, neg, clz etc. between two FragmentedArrays with different layouts, e.g. a WGMMA-layout accumulator plus a memory/distributed-layout operand, without an explicit to_layout.

Common situations: Feeding a value loaded in one layout into arithmetic with a dot-product result in WGMMA layout; mixing arrays from different primitives in custom Mosaic kernels; layouts differing only in subtle parameters (num_regs, row/col splits) so the mismatch is invisible in logs.

Related errors

• Output layout {out_layout} does not match the layout of the
• Pointwise operations on {bitwidth}-bit types are unsupported
• Incompatible FragmentedArray shapes
• Source and destination layouts aren't compatible for a broad
• All arrays must have the same layout, got {arr.layout} at in

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