jax-ml/jax · error · ValueError

Expected WGStridedFragLayout, got {arr.layout} at index {i}

Error message

Expected WGStridedFragLayout, got {arr.layout} at index {i}

What it means

In the WGStridedFragLayout branch, concatenate requires every array to also have a WGStridedFragLayout, because register concatenation relies on the shared strided structure. A TiledLayout (or other layout) mixed into the list raises ValueError with its index.

Source

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

        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}"
          )
        if arr.layout.vec_size != vec_size:
          raise ValueError(
              "All WGStridedFragLayout arrays must have the same vec_size,"
              f" got {arr.layout.vec_size} at index {i} (expected {vec_size})"
          )
      new_layout = WGStridedFragLayout(shape=new_shape, vec_size=vec_size)
      new_regs = np.concatenate([arr.registers for arr in arrays], axis=0)
      return FragmentedArray(
          _registers=new_regs, _layout=new_layout, _is_signed=arr0.is_signed
      )

    case WGSplatFragLayout():
      raise NotImplementedError(
          "Concatenating arrays with splat layout is not supported."
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert all fragments to WGStridedFragLayout (or all to TiledLayout) before concatenating, e.g., round-trip through memory with a consistent load path.
2. Group fragments by layout type and concatenate within each group before combining.
3. Check isinstance(arr.layout, WGStridedFragLayout) per element before the call.

Example fix

# before
out = FragmentedArray.concatenate([wgmma_acc, tiled_frag], axis=0)  # error at index 1
# after
tiled_as_strided = reload_with_strided_layout(tiled_frag)
out = FragmentedArray.concatenate([wgmma_acc, tiled_as_strided], axis=0)

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

assert all(isinstance(a.layout, WGStridedFragLayout) for a in arrays), [
    (i, type(a.layout).__name__) for i, a in enumerate(arrays)
]
out = FragmentedArray.concatenate(arrays, axis=0)

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

def all_strided(arrays) -> bool:
    return all(isinstance(a.layout, WGStridedFragLayout) for a in arrays)

try:
    out = FragmentedArray.concatenate(arrays, axis=0)
except ValueError as e:
    if 'WGStridedFragLayout' not in str(e): raise
    arrays = [to_strided(a) for a in arrays]  # normalize via memory round-trip
    out = FragmentedArray.concatenate(arrays, axis=0)

When it happens

Trigger: Concatenating a WGMMA accumulator fragment (strided layout) with a fragment loaded from tiled SMEM/GMEM (TiledLayout) in the same list.

Common situations: Mixing WGMMA outputs with directly-loaded fragments and concatenating them; pipeline refactors that change how one fragment is materialized; fragments created via different load helpers with different layout inference.

Related errors

• All arrays must have the same layout, got {arr.layout} at in
• dims and idxs must have the same length
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• cluster= must be at most 3D, got {self}.

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