jax-ml/jax · error · NotImplementedError

CollapseShapeOp with empty reassociation is not supported.

Error message

CollapseShapeOp with empty reassociation is not supported.

What it means

Mosaic GPU's layout inference cannot derive constraints for memref.CollapseShapeOp when the reassociation list is empty. An empty reassociation only occurs when collapsing a (1, ...) shape down to an empty (scalar-like) shape, a case the pass explicitly declines to handle.

Source

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

      cs.Equals(source_var, dest_var),
  ]
  return cs.ConstraintSystem(constraints=constraints), {
      source_var: [source],
      dest_var: [dest],
  }


@_add_constraint_system_derivation_rule(memref.CollapseShapeOp)
def _memref_collapse_shape_op_constraint_system(
    ctx: DerivationContext,
    op: memref.CollapseShapeOp,
) -> 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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard your code path: skip the collapse (or use the original memref) when the reassociation would be empty, e.g. when all dimensions are size 1
2. Instead of collapsing to an empty shape, collapse to shape [1] with reassociation [[0,...]] so reassociation is non-empty
3. Adjust tile sizes/block specs so shapes never degenerate to all-ones

Example fix

# before
collapsed = memref.collapse_shape(src, reassociation=[])

# after
if src.shape and all(d == 1 for d in src.shape):
    collapsed = memref.collapse_shape(src, reassociation=[[0] * len(src.shape)])
else:
    collapsed = memref.collapse_shape(src, reassociation=reassoc)

reassoc = build_reassociation(src.shape, dest.shape)
if not reassoc:  # collapsing (1,...) to ()
    skip_collapse = True

def is_collapse_supported(reassociation) -> bool:
    return len(reassociation) > 0

When it happens

Trigger: Calling memref.collapse_shape with reassociation=() on a memref whose shape is all 1s (e.g. (1,1) -> ()), producing an empty reassociation attribute inside a Mosaic GPU kernel.

Common situations: Generic reshape utilities that collapse all singleton dimensions; kernels parameterized so a tile shape degenerates to (1,...,1) at small sizes, triggering the empty-reassociation path.

Related errors

• Only unit strides are supported but got {op.static_strides}.
• CollapseShapeOp with non-contiguous strides is not supported
• Only a single tiling transform is supported when collapsing
• Transposed memrefs are not supported in ExpandShapeOp.
• Only scalar memrefs are supported, got {ref_shape}

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