jax-ml/jax · error · NotImplementedError

Only a single tiling transform is supported when collapsing

Error message

Only a single tiling transform is supported when collapsing a shape, but got {in_transforms=} and {out_transforms=}

What it means

collapse_shape lowering supports at most a single TileTransform on input and output; anything else (multiple transforms, non-tiling transforms) is unimplemented.

Source

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

# layout inference. It is not entirely clear what the best approach is.
def _check_collapse_shape(
    op: memref.CollapseShapeOp,
    in_transforms: Sequence[lc.MemRefTransform],
    out_transforms: Sequence[lc.MemRefTransform],
):
  if len(in_transforms) != len(out_transforms):
    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),
      {},
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure exactly one TileTransform in and one out, nothing else
2. Strip swizzle transforms before the collapse (or use untiled memrefs)
3. Restructure the kernel so collapse happens on untransformed data

def single_tile_transforms_only(in_ts, out_ts) -> bool:
    return (len(in_ts) == 1 and len(out_ts) == 1
            and isinstance(in_ts[0], lc.TileTransform)
            and isinstance(out_ts[0], lc.TileTransform))

When it happens

Trigger: memref.collapse_shape with more than one input transform, or transforms that aren't lc.TileTransform on either side.

Common situations: Collapsing a tensor that has both tiling and swizzle transforms attached; hitting the single-transform limitation of the collapse path.

Related errors

• Unsupported in/out transforms. In transform: {in_transform},
• CollapseShapeOp with empty reassociation is not supported.
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}

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