jax-ml/jax · error · ValueError

strides={ref_ty.get_strides_and_offset()[0]}, {ref_ty.shape=

Error message

strides={ref_ty.get_strides_and_offset()[0]}, {ref_ty.shape=}, {dim=}, {fold_rank=}

What it means

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.

Source

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

        f"Folding {fold_rank} dimensions starting from {dim} is out of bounds"
        f" for shape {new_shape}"
    )
  new_shape[dim : dim + fold_rank] = [
      math.prod(new_shape[dim : dim + fold_rank])
  ]
  identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(ref_ty.rank))
  contig_strided_1d = ir.Attribute.parse("strided<[1]>")
  # Not sure why but MLIR expects the strided 1D layout to disappear in this op.
  if ref_ty.layout == identity or ref_ty.layout == contig_strided_1d:
    new_layout = ir.AffineMapAttr.get(
        ir.AffineMap.get_identity(ref_ty.rank - fold_rank + 1)
    )
  elif _is_contiguous_shape_slice(ref_ty, slice(dim, dim + fold_rank)):
    new_strides, offset = ref_ty.get_strides_and_offset()
    new_strides[dim : dim + fold_rank] = [new_strides[dim + fold_rank - 1]]
    new_layout = ir.StridedLayoutAttr.get(offset, new_strides)
  else:
    raise ValueError(
        f"strides={ref_ty.get_strides_and_offset()[0]}, {ref_ty.shape=},"
        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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Make the region contiguous first: copy or materialize the data into a contiguous memref (e.g. via memref.copy into an alloca) before folding
  2. Fold a different, contiguous set of dimensions that matches the stride pattern
  3. Inspect ref_ty.get_strides_and_offset() (as the message prints) and adjust the layout or choose dims whose strides are nested multiples

Example fix

// before
folded = utils.memref_fold(sliced_ref, dim=1, fold_rank=2)  # sliced_ref has gaps
// after
contig = memref.alloca(ir.MemRefType.get(sliced_shape, elem_ty), [], [])
memref.copy(sliced_ref, contig)
folded = utils.memref_fold(contig, dim=1, fold_rank=2)
Defensive patterns

Strategy: validation

Validate before calling

ref_ty = ir.MemRefType(ref.type)
strides, _ = ref_ty.get_strides_and_offset()
# foldable iff slice is contiguous; quick check for the common row-major case:
def contiguous(dim, fold_rank, shape, strides):
    expected = 1
    for i in reversed(range(dim, dim + fold_rank)):
        if strides[i] != expected:
            return False
        expected *= shape[i]
    return True
assert contiguous(dim, fold_rank, ref_ty.shape, strides)

Prevention

When it happens

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

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

Related errors

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