jax-ml/jax · error · NotImplementedError

Concatenating arrays with strided layout is only supported a

Error message

Concatenating arrays with strided layout is only supported along axis 0

What it means

For WGStridedFragLayout fragments, concatenate is implemented by appending whole register arrays (row-major register blocks per warpgroup), which only preserves the strided layout semantics when concatenating along axis 0. Any other axis raises NotImplementedError.

Source

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

    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}"
          )
        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
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Concatenate along axis 0 instead (restructure the kernel so the joined dimension is the leading one).
2. Or go through memory: store fragments to a tiled/untiled reference, concatenate there, and reload with the strided layout.
3. Or convert fragments to a TiledLayout, concatenate on any axis, then convert back.

Example fix

# before
out = FragmentedArray.concatenate(accs, axis=1)  # WGStridedFragLayout -> error
# after
# transpose the split so joining happens along axis 0:
out = FragmentedArray.concatenate(accs, axis=0)

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

if isinstance(arrays[0].layout, WGStridedFragLayout) and axis != 0:
    axis = 0  # or restructure split so axis 0 is the joined dim
out = FragmentedArray.concatenate(arrays, axis=axis)

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

def can_concat_axis0_only(arrays) -> bool:
    return isinstance(arrays[0].layout, WGStridedFragLayout)

try:
    out = FragmentedArray.concatenate(arrays, axis=axis)
except NotImplementedError:
    # fall back: concat along axis 0, or via TiledLayout round-trip
    out = FragmentedArray.concatenate(arrays, axis=0)

When it happens

Trigger: Calling FragmentedArray.concatenate(arrs, axis=1) (or any nonzero axis) where arrays[0].layout is a WGStridedFragLayout (typical for WGMMA accumulator fragments).

Common situations: Reusing numpy-style concat code on WGMMA accumulator fragments; splitting accumulators along a non-leading dimension and trying to rejoin them; layout of intermediates changing to strided after a WGMMA op in a newer JAX.

Related errors

• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Cannot convert from {self.layout} to {new_layout}
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}

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