jax-ml/jax · error · NotImplementedError

Only 2D gathers/scatters for async load/store are supported.

Error message

Only 2D gathers/scatters for async load/store are supported.

What it means

Vector-indexed (gather/scatter) async loads and stores in Mosaic are only implemented for 2D global-memory tensors: if the source (async_load) or destination (async_store) memref does not have exactly 2 dimensions, a NotImplementedError is raised before constraints are built.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2228

) -> ConstraintSystemDerivationRuleResult:
  # We only support 2D gathers/scatters along the leading dimension. Tiling
  # either keeps the gather/scatter dimension leading or allows
  # collapsing leading dimensions to maintain contiguity without
  # transforming global memory.
  tiling_multiple = []
  for i, (size, index) in enumerate(zip(op.slice_lengths, op.indices, strict=True)):
    if size == -1:
      # This dimension does not appear in the final smem memref shape.
      continue
    if isinstance(index.type, ir.VectorType):
      if i != 0:
        raise NotImplementedError("Only leading gather dimensions allowed.")
      if isinstance(op, mgpu.AsyncStoreOp):
        gmem_shape = ir.MemRefType(op.destination.type).shape
      else:
        gmem_shape = ir.MemRefType(op.source.type).shape
      if len(gmem_shape) != 2:
        raise NotImplementedError("Only 2D gathers/scatters for async load/store are supported.")
      tiling_multiple.append(size)
      continue
    tiling_multiple.append(dynamic_gcd(size, index))

  operand_index = 1 if isinstance(op, mgpu.AsyncLoadOp) else 0
  operand = ValueSite(op, VariableType.OPERAND, operand_index)
  var = ctx.producer_ref(operand)
  constraints: list[cs.Constraint] = [
      cs.Divides(expr=var, tiling_multiple=tuple(tiling_multiple))
  ]
  if any(isinstance(idx.type, ir.VectorType) for idx in op.indices):
    element_bitwidth = utils.bitwidth(op.source.type.element_type)
    # This constraint enforces sufficient SMEM-alignment.
    # The transfer chunk needs to be 1024 bit-aligned. For each write in the
    # lowering we transfer 4 rows, so each row must be 256 bit-aligned.
    divisor = (1024 // 4) // element_bitwidth
    slice_lengths = [s for s in op.slice_lengths if s != -1]
    if slice_lengths and (slice_lengths[-1] % divisor):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape the global tensor to 2D before the gather (e.g. flatten leading dims into rows: [B, N, D] -> [B*N, D]) and adjust offsets accordingly
2. For 1D gathers, add a trivial second dimension of size 1 (memref.expand_shape / reshape to [N, 1] or [1, N]) so the rank is 2
3. For genuinely higher-rank access patterns, loop over the extra dims with 2D gathers per slice

Example fix

# before: src is memref<10000xf32> (rank 1) -> NotImplementedError
mgpu.async_load(src, smem, indices=vec_rows, slice_lengths=(1,))

# after: reshape to 2D first
src2 = reshape_to_2d(src)  # memref<10000x1xf32>
mgpu.async_load(src2, smem, indices=vec_rows, slice_lengths=(1, 1))

gmem = op.destination if isinstance(op, mgpu.AsyncStoreOp) else op.source
if any(isinstance(i.type, ir.VectorType) for i in op.indices) and len(ir.MemRefType(gmem.type).shape) != 2:
    gmem = reshape_to_2d(gmem)

def is_2d_gather_compatible(op) -> bool:
    gmem = op.destination if isinstance(op, mgpu.AsyncStoreOp) else op.source
    has_vec = any(isinstance(i.type, ir.VectorType) for i in op.indices)
    return (not has_vec) or len(ir.MemRefType(gmem.type).shape) == 2

When it happens

Trigger: Calling mgpu.async_load/async_store with a VectorType index where the gmem source/destination memref has rank != 2, e.g. a 1D tensor gather or a 3D batched gather.

Common situations: Gathering rows from a flattened 1D array, gathering from 3D+ activations in attention-style kernels, or reusing gather code written for 2D against higher-rank tensors after a refactor.

Related errors

• Only leading gather dimensions allowed.
• Unsupported TMA index shape {shape}
• Expected metadata to be 3-dimensional (M, K // 4, 2), but it
• copy_gmem_to_smem with a barrier is only supported Hopper an
• Splat layout does not support multimem

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