jax-ml/jax · error · ValueError

Source and destination layouts aren't compatible for a broad

Error message

Source and destination layouts aren't compatible for a broadcast

What it means

Thrown by FragmentedArray broadcast when the destination layout cannot be produced by reducing (removing) the newly added broadcast dimensions from the source layout. Mosaic GPU layouts must be structurally compatible across a broadcast: the non-broadcast dims of the result layout must exactly match the source layout with the new dims collapsed. If they don't match, broadcasting register fragments between layouts is not supported.

Source

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

      return FragmentedArray(
          _registers=np.tile(
              self.registers,
              math.prod(shape) // math.prod(self.shape),
          ),
          _layout=layout,
          _is_signed=self.is_signed,
      )
    if not isinstance(self.layout, TiledLayout) or not isinstance(layout, TiledLayout):
      raise NotImplementedError(self.layout, layout)
    if len(layout.base_tile_shape) != len(shape):
      raise NotImplementedError(
          "Tiling rank different than broadcast result rank, "
          f"{layout.base_tile_shape} vs {shape}"
      )
    new_dimensions = sorted(set(range(len(shape))) - set(source_dimensions))
    expected_layout = layout.reduce(new_dimensions)
    if expected_layout != self.layout:
      raise ValueError(
          "Source and destination layouts aren't compatible for a broadcast"
      )
    new_registers_shape = layout.registers_shape(shape)
    pre_broadcast_registers_shape = list(new_registers_shape)
    for new_dim in new_dimensions:
      for i, is_new in enumerate(layout.tiling.tile_dimension(new_dim)):
        if is_new:
          pre_broadcast_registers_shape[i] = 1
    # The broadcast for all dims but the vector_dim amounts to repeating the
    # registers along the new dimensions. Along the vector_dim, we actually need
    # to extend the vector length to change the type of the registers.
    if layout.vector_length != self.layout.vector_length:
      assert self.layout.vector_length == 1
      registers = np.empty_like(self.registers)
      for idx, reg in np.ndenumerate(self.registers):
        registers[idx] = utils.vector_concat([reg] * layout.vector_length)
    else:
      registers = self.registers

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Check the two layouts: assert layout.reduce(new_dims) == self.layout before broadcasting; align operand layouts with fg.to_layout or mgpu.repeat/lower operands to a common TiledLayout first.
Avoid broadcasting fragment-layout arrays directly; instead splat the scalar/operand explicitly (FragmentedArray.splat) or re-fragment the operand so both sides share one layout.
If the layouts look compatible, verify target_shape rank and that only new dims are being added — a mismatched existing dim also fails the equality check.

Example fix

# before
result = a.broadcast((m, n))  # a has WGStridedFragLayout, fails

# after
b = mgpu.fragmented_array(..., layout=TiledLayout(...))
a2 = a.to_layout(b.layout)
result = a2.broadcast((m, n))

Defensive patterns

Strategy: validation

Validate before calling

new_dims = sorted(set(range(len(target_shape))) - set(range(len(a.shape))))
assert a.layout.reduce(new_dims) == a.layout_expected, 'broadcast layouts incompatible'

Type guard

def can_broadcast(a, target_shape, src_dims):
    new_dims = sorted(set(range(len(target_shape))) - set(src_dims))
    try:
        return a.layout.reduce(new_dims) == a.layout
    except Exception:
        return False

Try / catch

try:
    out = a.broadcast(shape)
except ValueError as e:
    if 'compatible for a broadcast' in str(e):
        a = a.to_layout(common_layout)
        out = a.broadcast(shape)
    else:
        raise

Prevention

Keep both operands of pointwise ops on the same layout (convert with to_layout before combining).
Prefer TiledLayout for arrays that will be broadcast.
Unit-test kernel wrappers with all intended shapes to catch layout mismatches early.

When it happens

Trigger: Calling FragmentedArray.broadcast (or ops that lower to it, e.g. _pointwise with shape promotion, to_layout, __getitem__, custom primitive blocks) where target_shape adds dimensions whose reduced layout differs from self.layout — e.g. mixing WGStridedFragLayout/WGSplatFragLayout with a tiled result, or broadcasting to a shape whose tiling/vec dims don't line up with the source's.

Common situations: Writing Mosaic GPU kernels where an operand with a warpgroup fragment layout is broadcast against an array with a TiledLayout; manual layout construction via to_layout before a binary op; shape changes that add dims incompatible with the chosen layout.

Related errors

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