jax-ml/jax · error · ValueError

Can only infer one dimension

Error message

Can only infer one dimension

What it means

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.

Source

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

        f" {dim=}, {fold_rank=}"
    )

  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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace all but one None with concrete factors so only one dimension is inferred
2. Compute the missing factor yourself: inferred = size // prod(known factors) and pass integers only
3. If you truly want an ambiguous split, decide the factorization explicitly (e.g. [size//16, 16])

Example fix

# before
new_ref = utils.memref_unfold(ref, dim=1, factors=[None, None])
# after
new_ref = utils.memref_unfold(ref, dim=1, factors=[None, 16])

assert sum(f is None for f in factors) <= 1, 'only one inferred factor allowed'

def valid_unfold_factors(factors) -> bool:
    return sum(f is None for f in factors) <= 1

When it happens

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

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

Related errors

• Non-divisible unfold:
• Unexpected value for 'order' argument: {order}.
• 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

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