jax-ml/jax · error · ValueError

Folding {fold_rank} dimensions starting from {dim} is out of

Error message

Folding {fold_rank} dimensions starting from {dim} is out of bounds for shape {new_shape}

What it means

memref_fold tries to collapse fold_rank consecutive dimensions of a memref starting at axis `dim`, but the requested range extends past the memref's rank. JAX Mosaic raises this because the resulting folded shape cannot be constructed. It is a pure shape-arithmetic validation error.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:818


@overload
def memref_fold(ref: MultimemRef, dim, fold_rank) -> MultimemRef:
  ...


def memref_fold(
    ref: ir.Value | MultimemRef, dim, fold_rank
) -> ir.Value | MultimemRef:
  if isinstance(ref, MultimemRef):
    return MultimemRef(memref_fold(ref.ref, dim, fold_rank))

  ref_ty = ir.MemRefType(ref.type)
  new_shape = list(ref_ty.shape)
  if dim < 0:
    raise ValueError(f"Dimension {dim} is negative")
  if dim + fold_rank > len(new_shape):
    raise ValueError(
        f"Folding {fold_rank} dimensions starting from {dim} is out of bounds"
        f" for shape {new_shape}"
    )
  new_shape[dim : dim + fold_rank] = [
      math.prod(new_shape[dim : dim + fold_rank])
  ]
  identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))
  contig_strided_1d = ir.Attribute.parse("strided<[1]>")
  # Not sure why but MLIR expects the strided 1D layout to disappear in this op.
  if ref_ty.layout == identity or ref_ty.layout == contig_strided_1d:
    new_layout = ir.AffineMapAttr.get(
        ir.AffineMap.get_identity(ref_ty.rank - fold_rank + 1)
    )
  elif _is_contiguous_shape_slice(ref_ty, slice(dim, dim + fold_rank)):
    new_strides, offset = ref_ty.get_strides_and_offset()
    new_strides[dim : dim + fold_rank] = [new_strides[dim + fold_rank - 1]]
    new_layout = ir.StridedLayoutAttr.get(offset, new_strides)
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Check len(ir.MemRefType(ref.type).shape) and clamp: dim + fold_rank <= rank and dim >= 0 before calling
Recompute dim from the right-hand side of the shape if you intended to fold trailing dims: dim = rank - fold_rank
Reduce fold_rank so the range fits the memref rank

Example fix

// before
folded = utils.memref_fold(ref, dim=3, fold_rank=2)  # rank-3 ref
// after
rank = len(ir.MemRefType(ref.type).shape)
assert 0 <= dim and dim + fold_rank <= rank
folded = utils.memref_fold(ref, dim, fold_rank)

Defensive patterns

Strategy: validation

Validate before calling

ref_ty = ir.MemRefType(ref.type)
assert 0 <= dim and dim + fold_rank <= len(ref_ty.shape), (dim, fold_rank, ref_ty.shape)

Prevention

Compute fold ranges from len(ir.MemRefType(ref.type).shape), not from tensor sizes in another layout
Assert shape/rank invariants in debug builds of kernels

When it happens

Trigger: Calling memref_fold(ref, dim, fold_rank) (directly or via helpers that reshape layouts) where dim + fold_rank exceeds len(ref_ty.shape), e.g. folding 2 dims starting at dim=-1-handled rank on a rank-2 memref. Negative dim is also rejected just above.

Common situations: Writing a Mosaic GPU kernel and computing a fold range from loop variables or inferred shapes; off-by-one when dim is 0-based but computed from a size; folding more dims than remain after a previous unfold.

Related errors

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