{"record":{"id":"28bc39e8c8038b47","repo":"jax-ml/jax","slug":"can-only-infer-one-dimension","errorCode":null,"errorMessage":"Can only infer one dimension","messagePattern":"Can only infer one dimension","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":857,"sourceCode":"        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)\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):","sourceCodeStart":839,"sourceCodeEnd":875,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L839-L875","documentation":"memref_unfold splits one dimension into `factors` sub-dimensions, allowing at most one factor to be None (to be inferred from the existing size). Passing more than one None makes the split underdetermined, so JAX Mosaic raises this ValueError.","triggerScenarios":"memref_unfold(ref, dim, [None, None]) or any factors list/tuple with two or more None entries, often reached via _reshape/memref_reshape/memref_unsqueeze paths in a kernel.","commonSituations":"Writing a reshape helper that mirrors numpy-style -1 inference and passing multiple -1/None; migrating code from jax.numpy.reshape semantics where two -1s are also illegal but the error text differs.","solutions":["Replace all but one None with concrete factors so only one dimension is inferred","Compute the missing factor yourself: inferred = size // prod(known factors) and pass integers only","If you truly want an ambiguous split, decide the factorization explicitly (e.g. [size//16, 16])"],"exampleFix":"# before\nnew_ref = utils.memref_unfold(ref, dim=1, factors=[None, None])\n# after\nnew_ref = utils.memref_unfold(ref, dim=1, factors=[None, 16])","handlingStrategy":"validation","validationCode":"assert sum(f is None for f in factors) <= 1, 'only one inferred factor allowed'","typeGuard":"def valid_unfold_factors(factors) -> bool:\n    return sum(f is None for f in factors) <= 1","tryCatchPattern":null,"preventionTips":["Never port numpy reshape habits with multiple -1/None","Compute all but one factor explicitly from known tile shapes"],"tags":["jax","mosaic-gpu","memref","reshape","invalid-argument"],"backgroundTag":"reshape-multiple-inferred-dims","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}