jax-ml/jax · error · NotImplementedError

CollapseShapeOp with non-contiguous strides is not supported

Error message

CollapseShapeOp with non-contiguous strides is not supported.

What it means

Mosaic GPU's layout inference handles memref.CollapseShapeOp only when the source memref has contiguous row-major strides. If the actual strides differ from get_contiguous_strides(shape) (e.g. after a transpose or strided subview), the pass raises NotImplementedError because producing correct layout constraints for non-contiguous collapse would require additional checks.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2068

) -> ConstraintSystemDerivationRuleResult:
  reassociation = tuple(len(ir.ArrayAttr(idx)) for idx in op.reassociation)
  # This should only occur when going from a (1, ...) shape to an empty shape.
  # We can handle it if needed, but right now `CollapseShape` will not deal with
  # this case.
  if not reassociation:
    raise NotImplementedError(
        "CollapseShapeOp with empty reassociation is not supported."
    )

  source = ValueSite(op, VariableType.OPERAND, 0)
  source_var = ctx.producer_ref(source)
  dest = ValueSite(op, VariableType.RESULT, 0)
  dest_var = cs.Variable(dest)

  strides, _ = ir.MemRefType(source.value.type).get_strides_and_offset()
  # In this case, we'd need additional checks to produce a correct constraint.
  if strides != utils.get_contiguous_strides(source.shape):
    raise NotImplementedError(
        "CollapseShapeOp with non-contiguous strides is not supported."
    )

  # TODO(bchetioui): We could generate an inverse expression `ExpandShape` in
  # order to allow inferring layouts bidirectionally. This would allow removing
  # transforms from some kernels' BlockSpecs, but is not necessary at this time.
  collapse_expr = cs.CollapseShape(source_var, source.shape, reassociation)
  return cs.ConstraintSystem(constraints=[cs.Equals(dest_var, collapse_expr)]), {
      source_var: [source],
      dest_var: [dest],
  }


# `memref.load` and `memref.store` are used to load barrier phases which are
# scalars---the rule needn't do anything interesting, but we need to have it.
@_add_constraint_system_derivation_rule(memref.LoadOp)
@_add_constraint_system_derivation_rule(memref.StoreOp)
def _memref_load_store_op_constraint_system(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Materialize a contiguous copy of the memref (round-trip through registers or a fresh smem allocation) before collapse_shape
  2. Reorder operations: collapse first on the contiguous source, then apply transpose/strided access on the collapsed result
  3. If strides are contiguous but include a padding dimension, adjust the allocation/padding so get_strides_and_offset matches contiguous strides

Example fix

# before: src has strides [1, 512] but shape [512, 128] (padded/strided)
flat = memref.collapse_shape(src, reassociation=[[0, 1]])

# after: copy into contiguous buffer first
contig = alloc_contiguous_like(src)
copy(src, contig)
flat = memref.collapse_shape(contig, reassociation=[[0, 1]])
Defensive patterns

Strategy: validation

Validate before calling

strides, _ = ir.MemRefType(src.type).get_strides_and_offset()
if strides != utils.get_contiguous_strides(src.shape):
    src = copy_to_contiguous(src)  # materialize before collapse_shape

Type guard

def is_contiguous_memref(v) -> bool:
    m = ir.MemRefType(v.type)
    s, _ = m.get_strides_and_offset()
    return s == utils.get_contiguous_strides(m.shape)

Prevention

When it happens

Trigger: Applying memref.collapse_shape to a memref with non-contiguous strides: a transposed memref, a strided subview (e.g. a column slice), or a swizzled smem layout, inside a Mosaic GPU kernel.

Common situations: Flattening a tile after slicing columns (row stride > row length), reshaping transposed shared-memory tiles, or collapsing views created by TMA/smem transforms with padded or swizzled strides.

Related errors

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