jax-ml/jax · error · NotImplementedError

Collapsing the shape of a memref with non-contiguous strides

Error message

Collapsing the shape of a memref with non-contiguous strides is not supported

What it means

Collapsing a tiled memref is only implemented when the source memref is fully contiguous. Non-contiguous strides (e.g. from a prior subview or slice) make the tiling reduction invalid.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:2272

    raise ValueError(
        "Expected the same number of in/out transforms, but got "
        f"{in_transforms=} and {out_transforms=}"
    )
  if not in_transforms:
    return
  t_in, *in_transforms = in_transforms
  t_out, *_ = out_transforms
  if (in_transforms or
      not isinstance(t_in, lc.TileTransform) or
      not isinstance(t_out, lc.TileTransform)):
    raise NotImplementedError(
        "Only a single tiling transform is supported when collapsing a shape, "
        f"but got {in_transforms=} and {out_transforms=}"
    )
  src_ty = ir.MemRefType(op.src.type)
  strides, _ = src_ty.get_strides_and_offset()
  if strides != utils.get_contiguous_strides(src_ty.shape):
    raise NotImplementedError(
        "Collapsing the shape of a memref with non-contiguous strides is not "
        "supported"
    )
  reassociation = tuple(len(ir.ArrayAttr(idx)) for idx in op.reassociation)

  collapsed_tiling = cs.reduce_expression(
      cs.CollapseShape(cs.SMEMTransforms(t_in, None), tuple(src_ty.shape),
                       reassociation),
      {},
  )

  if isinstance(collapsed_tiling, cs.Unsatisfiable):
    raise ValueError(f"Input tiling {t_in.tiling} is not compatible with {op}")

  assert isinstance(collapsed_tiling, cs.SMEMTransforms)
  expected_t_out = collapsed_tiling.tiling
  assert expected_t_out is not None
  if expected_t_out != t_out:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Copy the strided data into a fresh contiguous buffer before collapsing
2. Collapse before taking subviews/slices
3. Ensure the allocation has no padding so strides are contiguous

Example fix

// before
collapsed = t.memref.collapse_shape(sliced_strided_ref, reassoc)
// after
contig = t.copy_to_contiguous(sliced_strided_ref)  # or store/load
collapsed = t.memref.collapse_shape(contig, reassoc)

strides, _ = ir.MemRefType(src.type).get_strides_and_offset()
assert strides == utils.get_contiguous_strides(ir.MemRefType(src.type).shape)

def is_contiguous(ty) -> bool:
    strides, _ = ty.get_strides_and_offset()
    return strides == utils.get_contiguous_strides(ty.shape)

When it happens

Trigger: memref.collapse_shape on a tiled memref whose src_ty.get_strides_and_offset() != contiguous strides for its shape, e.g. collapsing a sliced or padded buffer.

Common situations: Chaining slice/subview (which produces strided views) followed by collapse_shape on smem tiles; or collapsing a padded allocation.

Related errors

• CollapseShapeOp with non-contiguous strides is not supported
• Data type bitcast is only supported for contiguous 1D arrays
• Expected the same number of in/out transforms, but got {in_t
• Only a single tiling transform is supported when collapsing
• Input tiling {t_in.tiling} is not compatible with {op}

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