jax-ml/jax · error · ValueError

Cannot broadcast shape {self.shape} to layout {o.layout}

Error message

Cannot broadcast shape {self.shape} to layout {o.layout}

What it means

In _pointwise, when the other operand is a splat (WGSplatFragLayout), the library tries to materialize it at self's shape via WGSplatFragLayout.can_broadcast_to. If self.shape is not broadcast-compatible with the splat's layout shape, broadcasting is impossible and ValueError is raised.

Source

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

        return self.broadcast(output_shape)._pointwise(
            op, *other, output_is_signed=output_is_signed
        )

    other_arrs = []
    for o in other:
      if not isinstance(o, FragmentedArray):
        if isinstance(o, (float, int)):
          o = utils.c(o, self.mlir_dtype)
        elif not isinstance(o, ir.Value):
          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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Recreate the splat with a shape that broadcasts: FragmentedArray.splat(value, shape=self.shape, ...) or shape=()
Reshape/expand self so the dims align with the splat's layout shape
Verify tile-shape consistency between constants and operands before fusing ops

Example fix

# before
bias = FragmentedArray.splat(b, shape=(1, 64), layout=L)
z = x + bias  # x.shape == (64,)
# after
bias = FragmentedArray.splat(b, shape=x.shape, layout=x.layout)
z = x + bias

Defensive patterns

Strategy: validation

Validate before calling

if isinstance(o.layout, WGSplatFragLayout):
    assert o.layout.can_broadcast_to(x.shape), f"{o.layout} cannot broadcast to {x.shape}"

Try / catch

try:
    z = x + o
except ValueError as e:
    if 'Cannot broadcast' in str(e):
        o = FragmentedArray.splat(o.registers.flat[0], shape=x.shape, layout=x.layout)
        z = x + o
    else:
        raise

Prevention

Build splat constants with shape=x.shape (or ()) of the operand they combine with
Keep constant tile shapes in one config shared by the whole kernel

When it happens

Trigger: Combining a FragmentedArray with a splat constant whose stored layout shape doesn't broadcast to self.shape, e.g. adding a splat of shape (8, 16) to an array of shape (8,) or (16, 8) with mismatched leading dims, or shape (3,) against layouts implying power-of-two shapes.

Common situations: Adding bias/scale splats built with a different tile shape than the accumulation array; reusing constants across kernels with different tiling; constructing splats from scalars with an explicit stale shape argument.

Related errors

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