jax-ml/jax · error · ValueError

All arrays must have the same layout, got {arr.layout} at in

Error message

All arrays must have the same layout, got {arr.layout} at index {i} (expected {arr0.layout})

What it means

In the TiledLayout branch, concatenate verifies every array has the identical layout because np.concatenate on registers is only layout-preserving when tiles, register arrangement, and swizzle all match. A differing TiledLayout raises ValueError with both layouts printed.

Source

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

      raise ValueError(
          f"All arrays must have the same signedness, got {arr.is_signed} at"
          f" index {i} (expected {arr0.is_signed})"
      )
    for d in range(rank):
      if d != axis and arr.shape[d] != arr0.shape[d]:
        raise ValueError(
            "All arrays must have matching shapes along non-concatenated"
            f" dimensions, got shape {arr.shape} at index {i} (expected dim"
            f" {d} to be {arr0.shape[d]})"
        )
    new_shape[axis] += arr.shape[axis]
  new_shape = tuple(new_shape)

  match arr0.layout:
    case TiledLayout():
      for i, arr in enumerate(arrays[1:], start=1):
        if arr.layout != arr0.layout:
          raise ValueError(
              f"All arrays must have the same layout, got {arr.layout} at"
              f" index {i} (expected {arr0.layout})"
          )
      new_regs = np.concatenate([arr.registers for arr in arrays], axis=axis)
      return FragmentedArray(
          _registers=new_regs, _layout=arr0.layout, _is_signed=arr0.is_signed
      )

    case WGStridedFragLayout(vec_size=vec_size):
      if axis != 0:
        raise NotImplementedError(
            "Concatenating arrays with strided layout is only supported along"
            " axis 0"
        )
      for i, arr in enumerate(arrays[1:], start=1):
        if not isinstance(arr.layout, WGStridedFragLayout):
          raise ValueError(
              f"Expected WGStridedFragLayout, got {arr.layout} at index {i}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Re-materialize the mismatched fragments through the same layout as arrays[0] (store to GMEM/SMEM and reload with arrays[0].layout parameters).
2. Ensure all fragments are created by the same load/config path so their TiledLayout compares equal.
3. Print and compare .layout of each array before concat to find the divergent one early.

Example fix

# before
out = FragmentedArray.concatenate([a, b], axis=0)  # b.layout != a.layout
# after
b_relaid = b.to_tiled_layout_like(a)  # store + reload with a's tiling params
out = FragmentedArray.concatenate([a, b_relaid], axis=0)

lay = arrays[0].layout
bad = [i for i, a in enumerate(arrays) if a.layout != lay]
assert not bad, f'layout mismatch at {bad}: {lay}'
out = FragmentedArray.concatenate(arrays, axis=axis)

from jax.experimental.mosaic.gpu.fragmented_array import TiledLayout

def all_same_tiled_layout(arrays) -> bool:
    return isinstance(arrays[0].layout, TiledLayout) and all(
        a.layout == arrays[0].layout for a in arrays
    )

When it happens

Trigger: Concatenating fragments built with different tiled layouts, e.g. different tile shapes, different register orderings, or fragments produced by load_tiled with different swizzle values.

Common situations: Producing fragments in different passes or with different tiling parameters (tile shape, swizzle) and then combining them; JAX version changes altering default tiled layouts; mixing fragments from GMEM loads with fragments constructed directly from registers.

Related errors

• Expected WGStridedFragLayout, got {arr.layout} at index {i}
• dims and idxs must have the same length
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• cluster= must be at most 3D, got {self}.
• grid_names must have the same length as grid, got {self}.

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