jax-ml/jax · error · NotImplementedError

Arrays with the splat layout can only be stored when they ha

Error message

Arrays with the splat layout can only be stored when they have a single element or a multiple of {WARPGROUP_SIZE} elements

What it means

Storing an array with the splat layout requires the total element count to be either a single element or an exact multiple of WARPGROUP_SIZE (128) — the hardware stores splat fragments in warpgroup-wide vector chunks of vec_size = 64/bitwidth elements, so other sizes cannot be lowered.

Source

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

    ref = utils.memref_reshape(ref, (*(1 for _ in ref_ty.shape), *ref_ty.shape))
    return cls.load_tiled(
        ref, swizzle=swizzle, is_signed=is_signed, layout=layout, optimized=optimized
    )

  def _store_untiled_splat(self, ref: ir.Value):
    if math.prod(self.shape) == 1:
      c0 = c(0, ir.IndexType.get())
      memref.store(
          self.registers.flat[0], ref, [c0] * len(ir.MemRefType(ref.type).shape)
      )
      return

    vec_size = 64 // mgpu.bitwidth(self.mlir_dtype)
    if np.prod(self.shape) < vec_size * WARPGROUP_SIZE:
      vec_size = 1

    if np.prod(self.shape) % WARPGROUP_SIZE * vec_size:
      raise NotImplementedError(
          "Arrays with the splat layout can only be stored when they have a"
          f" single element or a multiple of {WARPGROUP_SIZE} elements"
      )

    fa = FragmentedArray.splat(
        self.registers.flat[0],
        self.shape,
        layout=WGStridedFragLayout(shape=self.shape, vec_size=vec_size),
        is_signed=self.is_signed,
    )
    fa.store_untiled(ref)

  def store_tiled_async(
      self,
      ref: ir.Value,
      barrier: utils.BarrierRef,
      cluster_dim: gpu.Dimension,
      cluster_idx: ir.Value,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pick block/shape sizes that are multiples of WARPGROUP_SIZE (128) when a splat-layout array will be stored.
If the value is a true scalar, store a 1-element splat instead of splatting to the full shape.
Convert the array to a TiledLayout before storing if the shape can't be changed (fa.to_layout(TiledLayout(...))).store(ref)).

Example fix

# before
fa = FragmentedArray.splat(c)  # shape (100,)
fa.store(ref)  # 100 not a multiple of 128

# after
# choose warpgroup-aligned block
fa = FragmentedArray.splat(c)  # shape (128,)
fa.store(ref)

Defensive patterns

Strategy: validation

Validate before calling

import numpy as np
from jax.experimental.mosaic.gpu.fragmented_array import WGSplatFragLayout, WARPGROUP_SIZE
if isinstance(fa.layout, WGSplatFragLayout):
    n = np.prod(fa.shape)
    assert n == 1 or (n * (64 // mgpu.bitwidth(fa.mlir_dtype))) % WARPGROUP_SIZE == 0, 'splat store needs 1 or warpgroup-multiple elements'

Type guard

def splat_storeable(fa):
    n = int(np.prod(fa.shape))
    return n == 1 or n % WARPGROUP_SIZE == 0

Prevention

Choose block sizes as multiples of 128 (WARPGROUP_SIZE) when splat arrays are stored.
Store single-element splats for true scalars instead of splatting to full shapes.
Convert to TiledLayout via to_layout when the shape can't be aligned.

When it happens

Trigger: FragmentedArray._store_untiled_splat / store on a WGSplatFragLayout array whose np.prod(shape) is not 1 and not a multiple of 128 (as combined with vec_size in the modulo check) — e.g. storing a splatted array of 3, 100, or 1000 elements.

Common situations: Splatting a scalar constant over an arbitrary-shaped result then storing it; block sizes not chosen as multiples of 128 (one warpgroup); changing tile shapes during kernel tuning so a splat operand's size is no longer warpgroup-aligned.

Related errors

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