{"record":{"id":"2ffd29958f422a68","repo":"jax-ml/jax","slug":"non-divisible-unfold","errorCode":null,"errorMessage":"Non-divisible unfold:","messagePattern":"Non-divisible unfold:","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":860,"sourceCode":"  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)\n  if sum(f is None for f in factors) > 1:\n    raise ValueError(\"Can only infer one dimension\")\n  known_factor_prod = np.prod([f for f in factors if f is not None])\n  if new_shape[dim] % known_factor_prod:\n    raise ValueError(\"Non-divisible unfold:\", new_shape[dim], factors)\n  factors = tuple(\n      new_shape[dim] // known_factor_prod if f is None else f for f in factors\n  )\n  new_shape[dim : dim + 1] = factors\n  identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))\n  contig_strided_1d = ir.Attribute.parse(\"strided<[1]>\")\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 + len(factors) - 1)\n    )\n  else:\n    new_strides, offset = ref_ty.get_strides_and_offset()\n    prev_stride = new_strides[dim]\n    inserted_strides = []\n    for f in reversed(factors):\n      inserted_strides.append(prev_stride)\n      prev_stride *= f\n    new_strides[dim : dim + 1] = reversed(inserted_strides)","sourceCodeStart":842,"sourceCodeEnd":878,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L842-L878","documentation":"memref_unfold requires the product of the known (non-None) factors to evenly divide the current dimension size. If new_shape[dim] % prod(factors) != 0, the split cannot produce integer sub-dimensions and Mosaic raises this ValueError with the dim size and factors.","triggerScenarios":"memref_unfold(ref, dim, [4, 8]) on a dim of size 30; or factors [None, 5] where the dim size is not a multiple of 5. Reached through _reshape/memref_reshape/memref_unsqueeze in kernel code.","commonSituations":"Hardcoding tile shapes (e.g. 128, 16) that don't divide the tensor extent; changing tensor shapes in an experiment without updating unfold factors; off-by-one in the dim index picking the wrong axis size.","solutions":["Check size = ir.MemRefType(ref.type).shape[dim] and choose factors whose product divides size","Pad or slice the tensor so the dim is divisible before unfolding","If using an inferred factor (None), ensure the known factors divide the dim; otherwise supply the full factorization"],"exampleFix":"# before\nref2 = utils.memref_unfold(ref, dim=1, factors=[16, 16])  # dim size 100\n# after\nsize = ir.MemRefType(ref.type).shape[1]\nassert size % 256 == 0, f'dim size {size} not divisible'\nref2 = utils.memref_unfold(ref, dim=1, factors=[16, 16])","handlingStrategy":"validation","validationCode":"size = ir.MemRefType(ref.type).shape[dim]\nknown = math.prod(f for f in factors if f is not None)\nassert size % known == 0, (size, factors)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive tile factors from the dim size itself (e.g. size // 16, 16)","Pad tensors to multiples of your tile shape before unfolding"],"tags":["jax","mosaic-gpu","memref","reshape","non-divisible-shape"],"backgroundTag":"reshape-dimension-not-divisible","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}