jax-ml/jax · error · ValueError

All WGStridedFragLayout arrays must have the same vec_size,

Error message

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

What it means

Mosaic GPU's FragmentedArray.concat requires all concatenated arrays that use WGStridedFragLayout to share the same vec_size (vector width). The loop validates each array's layout against the first array's vec_size and raises ValueError on mismatch. This invariant is required because a single WGStridedFragLayout with one vec_size is constructed for the concatenated result.

Source

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

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

    case layout:
      assert_never(layout)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make all inputs use the same vec_size: convert/reshape fragments (e.g. via layout conversion ops) so every WGStridedFragLayout array matches the first array's vec_size before concat
Check arr.layout.vec_size for every input before calling concat and re-emit the ones that differ
If mixing widths is intentional, concatenate via materializing to registers/tensors and re-fragmenting instead of FragmentedArray.concat

Example fix

# before
result = FragmentedArray.concat([a, b])  # a.vec_size=2, b.vec_size=4 -> ValueError

# after
assert all(arr.layout.vec_size == arrays[0].layout.vec_size for arr in arrays)
result = FragmentedArray.concat([a, b])  # both vec_size=4

Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental.mosaic.gpu import fragmented_array as fa

def can_concat(arrays):
  first = arrays[0].layout
  if isinstance(first, fa.WGStridedFragLayout):
    return all(
        isinstance(a.layout, fa.WGStridedFragLayout)
        and a.layout.vec_size == first.vec_size
        for a in arrays
    )
  return True

Type guard

def same_vec_size(arrays) -> bool:
  vs = getattr(arrays[0].layout, 'vec_size', None)
  return vs is None or all(getattr(a.layout, 'vec_size', None) == vs for a in arrays)

Try / catch

try:
  out = FragmentedArray.concat(arrays)
except ValueError as e:
  if 'vec_size' in str(e):
    raise RuntimeError(f'Incompatible vec_size: {[a.layout for a in arrays]}') from e
  raise

Prevention

Assert vec_size equality across fragment arrays before concat
Avoid mixing fragments produced under different vectorization settings in one kernel
Log arr.layout for each input when building kernels that concatenate fragments

When it happens

Trigger: Calling FragmentedArray.concat (or APIs building on it) where the first array has a WGStridedFragLayout with vec_size=N but a later array (index i) has vec_size != N, e.g. mixing arrays produced with different element vectorization (vec2 vs vec4 register fragments).

Common situations: Combining tensors materialized with different wgmma/mma vector widths, upgrading/downgrading Mosaic versions that changed default vec_size, or manually constructing FragmentedArray instances with mismatched layouts before concatenation.

Related errors

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