jax-ml/jax · error · NotImplementedError

Only leading gather dimensions allowed.

Error message

Only leading gather dimensions allowed.

What it means

In async gather/scatter (vector-indexed async load/store), Mosaic only supports gather dimensions that appear first: if a vector-typed index occurs at any position other than i == 0, layout inference raises NotImplementedError, because mixed leading-static/gather dims would need more general constraint modeling.

Source

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

@_add_constraint_system_derivation_rule(mgpu.AsyncLoadOp)
@_add_constraint_system_derivation_rule(mgpu.AsyncStoreOp)
def _async_load_store_constraint_system(
    ctx: DerivationContext,
    op: mgpu.AsyncLoadOp | mgpu.AsyncStoreOp,
) -> 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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose the global-memory tensor (or its layout) so the gathered dimension becomes dimension 0, then gather along rows
2. Reorder the indices tuple so the vector index is the first operand and preceding dims use size == -1 (sliced-away dims)
3. Use non-async per-element load/store for non-leading gather dims if performance permits

Example fix

# before: gather along dim 1 (i == 1 with VectorType) -> NotImplementedError
mgpu.async_load(src_T, smem, indices=(scalar_row, vec_cols), slice_lengths=(1, 1))

# after: transpose so gather is along dim 0
src = transpose(src_T)
mgpu.async_load(src, smem, indices=(vec_cols, scalar_row_or_none), slice_lengths=(1, 1))

for i, (size, idx) in enumerate(zip(slice_lengths, indices)):
    if size != -1 and not isinstance(idx.type, ir.VectorType):
        continue
    if size != -1 and isinstance(idx.type, ir.VectorType) and i != 0:
        raise ValueError('transpose tensors so gather dim is dim 0')

def gather_dims_are_leading(slice_lengths, indices) -> bool:
    return all(
        size == -1 or not isinstance(idx.type, ir.VectorType) or i == 0
        for i, (size, idx) in enumerate(zip(slice_lengths, indices))
    )

When it happens

Trigger: Constructing mgpu.async_load or mgpu.async_store where slice_lengths[i] != -1 with a VectorType indices[i] at position i > 0 — e.g. gathering along the second (column) dimension instead of the first (row) dimension of a 2D tensor.

Common situations: Column gathers (indexing the fast/last axis with vectors), or building indices tuples where a static scalar index precedes the vector index; porting Triton gather code that gathers along axis 1.

Related errors

• Only 2D gathers/scatters for async load/store are supported.
• Unsupported TMA index shape {shape}
• copy_gmem_to_smem with a barrier is only supported Hopper an
• Splat layout does not support multimem
• async_copy requires all GMEM strides except the last one to

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