{"record":{"id":"920627c194c8244e","repo":"jax-ml/jax","slug":"strides-ref-ty-get-strides-and-offset-0-ref","errorCode":null,"errorMessage":"strides={ref_ty.get_strides_and_offset()[0]}, {ref_ty.shape=}, {dim=}, {fold_rank=}","messagePattern":"strides=(.+?), (.+?), (.+?), (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":837,"sourceCode":"        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:\n    raise ValueError(\n        f\"strides={ref_ty.get_strides_and_offset()[0]}, {ref_ty.shape=},\"\n        f\" {dim=}, {fold_rank=}\"\n    )\n\n  new_ty = ir.MemRefType.get(\n      new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space\n  )\n  assoc = [[d] for d in range(dim)]\n  assoc.append([dim + i for i in range(fold_rank)])\n  assoc.extend([d] for d in range(dim + fold_rank, ref_ty.rank))\n  assert len(assoc) == new_ty.rank\n  return memref.collapse_shape(new_ty, ref, assoc)\n\n\ndef memref_unfold(ref: ir.Value, dim, factors) -> ir.Value:\n  \"\"\"Unfolds dim into two dimensions, the size of leading one given be major_factor.\"\"\"\n  ref_ty = ir.MemRefType(ref.type)\n  new_shape = list(ref_ty.shape)","sourceCodeStart":819,"sourceCodeEnd":855,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L819-L855","documentation":"memref_fold can only merge dimensions whose memory layout is compatible: either the folded slice is contiguous, or the special layout case above applies. If the strides of dims [dim, dim+fold_rank) are neither contiguous nor the foldable pattern, JAX Mosaic refuses to build the new strided layout. This protects against silently producing a memref whose linearization no longer matches memory.","triggerScenarios":"Calling memref_fold on a memref with a strided/non-contiguous layout (e.g. a slice of a larger buffer, a transposed view, or views with padded strides) where new_shape folding succeeds but _is_contiguous_shape_slice and the preceding layout branch both fail.","commonSituations":"Folding dims of a memref produced by memref_slice, memref_reinterpret_cast, or TMA/async-copy views with non-unit leading strides; feeding an arbitrary layout from a lower-level MLIR builder.","solutions":["Make the region contiguous first: copy or materialize the data into a contiguous memref (e.g. via memref.copy into an alloca) before folding","Fold a different, contiguous set of dimensions that matches the stride pattern","Inspect ref_ty.get_strides_and_offset() (as the message prints) and adjust the layout or choose dims whose strides are nested multiples"],"exampleFix":"// before\nfolded = utils.memref_fold(sliced_ref, dim=1, fold_rank=2)  # sliced_ref has gaps\n// after\ncontig = memref.alloca(ir.MemRefType.get(sliced_shape, elem_ty), [], [])\nmemref.copy(sliced_ref, contig)\nfolded = utils.memref_fold(contig, dim=1, fold_rank=2)","handlingStrategy":"validation","validationCode":"ref_ty = ir.MemRefType(ref.type)\nstrides, _ = ref_ty.get_strides_and_offset()\n# foldable iff slice is contiguous; quick check for the common row-major case:\ndef contiguous(dim, fold_rank, shape, strides):\n    expected = 1\n    for i in reversed(range(dim, dim + fold_rank)):\n        if strides[i] != expected:\n            return False\n        expected *= shape[i]\n    return True\nassert contiguous(dim, fold_rank, ref_ty.shape, strides)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Materialize slices into fresh contiguous allocas before folding dims","Print get_strides_and_offset() when a fold fails to spot non-contiguous views"],"tags":["jax","mosaic-gpu","memref","strides","layout"],"backgroundTag":"non-contiguous-tensor-view","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}