jax-ml/jax · error · NotImplementedError

Concatenating arrays with splat layout is not supported.

Error message

Concatenating arrays with splat layout is not supported.

What it means

FragmentedArray.concat handles strided-fragment layouts but explicitly does not support WGSplatFragLayout (a broadcast/splat register layout). The match statement raises NotImplementedError as a guard because concatenating splatted registers has no defined semantics in this code path. Developers must materialize or reshape the splat array first.

Source

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

        )
      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

1. Materialize the splat array into a strided/regular fragment layout (e.g. broadcast its registers into full storage) before concat
2. Compute the concatenated result another way: pad or reshape operands so splat values are produced post-concatenation
3. File/upstream a feature request if concat-of-splat is needed, since this is NotImplementedError by design

Example fix

# before
out = FragmentedArray.concat([splat_arr, other])  # splat_arr has WGSplatFragLayout

# after
materialized = splat_arr # convert to a strided fragment layout first
out = FragmentedArray.concat([materialized, other])

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

def concat_safe(arrays):
  if any(isinstance(a.layout, fa.WGSplatFragLayout) for a in arrays):
    raise TypeError('Splat-layout arrays cannot be concatenated; materialize them first')
  return fa.FragmentedArray.concat(arrays)

def is_splat(arr) -> bool:
  from jax.experimental.mosaic.gpu import fragmented_array as fa
  return isinstance(arr.layout, fa.WGSplatFragLayout)

try:
  out = FragmentedArray.concat(arrays)
except NotImplementedError as e:
  if 'splat' in str(e):
    arrays = [materialize(a) for a in arrays]  # your materialization
    out = FragmentedArray.concat(arrays)
  else:
    raise

When it happens

Trigger: Calling FragmentedArray.concat where any (here the first, arr0) array has WGSplatFragLayout, e.g. concatenating an array created from a scalar broadcast with normal strided-fragment arrays.

Common situations: Broadcasting a scalar constant into a kernel and then concatenating it with computed fragments; using older Mosaic versions where splat arrays silently worked differently after layout refactors.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• All WGStridedFragLayout arrays must have the same vec_size,
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

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