jax-ml/jax · error · ValueError

Non-divisible unfold:

Error message

Non-divisible unfold:

What it means

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.

Source

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

  new_ty = ir.MemRefType.get(
      new_shape, ref_ty.element_type, new_layout, ref_ty.memory_space
  )
  assoc = [[d] for d in range(dim)]
  assoc.append([dim + i for i in range(fold_rank)])
  assoc.extend([d] for d in range(dim + fold_rank, ref_ty.rank))
  assert len(assoc) == new_ty.rank
  return memref.collapse_shape(new_ty, ref, assoc)


def memref_unfold(ref: ir.Value, dim, factors) -> ir.Value:
  """Unfolds dim into two dimensions, the size of leading one given be major_factor."""
  ref_ty = ir.MemRefType(ref.type)
  new_shape = list(ref_ty.shape)
  if sum(f is None for f in factors) > 1:
    raise ValueError("Can only infer one dimension")
  known_factor_prod = np.prod([f for f in factors if f is not None])
  if new_shape[dim] % known_factor_prod:
    raise ValueError("Non-divisible unfold:", new_shape[dim], factors)
  factors = tuple(
      new_shape[dim] // known_factor_prod if f is None else f for f in factors
  )
  new_shape[dim : dim + 1] = factors
  identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))
  contig_strided_1d = ir.Attribute.parse("strided<[1]>")
  if ref_ty.layout == identity or ref_ty.layout == contig_strided_1d:
    new_layout = ir.AffineMapAttr.get(
        ir.AffineMap.get_identity(ref_ty.rank + len(factors) - 1)
    )
  else:
    new_strides, offset = ref_ty.get_strides_and_offset()
    prev_stride = new_strides[dim]
    inserted_strides = []
    for f in reversed(factors):
      inserted_strides.append(prev_stride)
      prev_stride *= f
    new_strides[dim : dim + 1] = reversed(inserted_strides)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check size = ir.MemRefType(ref.type).shape[dim] and choose factors whose product divides size
2. Pad or slice the tensor so the dim is divisible before unfolding
3. If using an inferred factor (None), ensure the known factors divide the dim; otherwise supply the full factorization

Example fix

# before
ref2 = utils.memref_unfold(ref, dim=1, factors=[16, 16])  # dim size 100
# after
size = ir.MemRefType(ref.type).shape[1]
assert size % 256 == 0, f'dim size {size} not divisible'
ref2 = utils.memref_unfold(ref, dim=1, factors=[16, 16])

size = ir.MemRefType(ref.type).shape[dim]
known = math.prod(f for f in factors if f is not None)
assert size % known == 0, (size, factors)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Can only infer one dimension
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unfolding dimensions is not supported when commuting an `Un
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w

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