{"record":{"id":"0ea484b5c46b3d03","repo":"jax-ml/jax","slug":"folding-fold-rank-dimensions-starting-from-dim","errorCode":null,"errorMessage":"Folding {fold_rank} dimensions starting from {dim} is out of bounds for shape {new_shape}","messagePattern":"Folding (.+?) dimensions starting from (.+?) is out of bounds for shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":818,"sourceCode":"\n\n@overload\ndef memref_fold(ref: MultimemRef, dim, fold_rank) -> MultimemRef:\n  ...\n\n\ndef memref_fold(\n    ref: ir.Value | MultimemRef, dim, fold_rank\n) -> ir.Value | MultimemRef:\n  if isinstance(ref, MultimemRef):\n    return MultimemRef(memref_fold(ref.ref, dim, fold_rank))\n\n  ref_ty = ir.MemRefType(ref.type)\n  new_shape = list(ref_ty.shape)\n  if dim < 0:\n    raise ValueError(f\"Dimension {dim} is negative\")\n  if dim + fold_rank > len(new_shape):\n    raise ValueError(\n        f\"Folding {fold_rank} dimensions starting from {dim} is out of bounds\"\n        f\" for shape {new_shape}\"\n    )\n  new_shape[dim : dim + fold_rank] = [\n      math.prod(new_shape[dim : dim + fold_rank])\n  ]\n  identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))\n  contig_strided_1d = ir.Attribute.parse(\"strided<[1]>\")\n  # Not sure why but MLIR expects the strided 1D layout to disappear in this op.\n  if ref_ty.layout == identity or ref_ty.layout == contig_strided_1d:\n    new_layout = ir.AffineMapAttr.get(\n        ir.AffineMap.get_identity(ref_ty.rank - fold_rank + 1)\n    )\n  elif _is_contiguous_shape_slice(ref_ty, slice(dim, dim + fold_rank)):\n    new_strides, offset = ref_ty.get_strides_and_offset()\n    new_strides[dim : dim + fold_rank] = [new_strides[dim + fold_rank - 1]]\n    new_layout = ir.StridedLayoutAttr.get(offset, new_strides)\n  else:","sourceCodeStart":800,"sourceCodeEnd":836,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L800-L836","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nfolded = utils.memref_fold(ref, dim=3, fold_rank=2)  # rank-3 ref\n// after\nrank = len(ir.MemRefType(ref.type).shape)\nassert 0 <= dim and dim + fold_rank <= rank\nfolded = utils.memref_fold(ref, dim, fold_rank)","handlingStrategy":"validation","validationCode":"ref_ty = ir.MemRefType(ref.type)\nassert 0 <= dim and dim + fold_rank <= len(ref_ty.shape), (dim, fold_rank, ref_ty.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","mosaic-gpu","memref","shape-mismatch","index-out-of-bounds"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}