jax-ml/jax · error · ValueError

Input tiling {t_in.tiling} is not compatible with {op}

Error message

Input tiling {t_in.tiling} is not compatible with {op}

What it means

The input tiling is not collapsible under the given reassociation: the symbolic collapse computation (cs.CollapseShape) returned Unsatisfiable, meaning no output tiling exists that preserves the tile semantics for this reassociation pattern.

Source

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

        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:
    raise ValueError(
        "Input/output tiling mismatch when attempting to collapse a shape. "
        f"Expected output tiling to be {expected_t_out.tiling} for input "
        f"tiling {t_in.tiling}, but got {t_out.tiling}"
    )


@_register_lowering(memref.CollapseShapeOp, support_warp_semantics=True)
def _memref_collapse_shape_op_lowering_rule(
    ctx: LoweringContext, op: memref.CollapseShapeOp
) -> Sequence[ir.Value]:
  del ctx

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change the reassociation so tile boundaries are respected (collapse whole tiles together)
2. Adjust tile shape so tiling aligns with the collapse groups
3. Retile or untile before collapsing

try:
    lowered = lower_collapse(op)
except ValueError as e:
    if 'not compatible' in str(e):
        regroup_reassociation(op)  # or untile first
    raise

When it happens

Trigger: memref.collapse_shape where a tile group would merge dimensions in a way incompatible with the tiling, e.g. collapsing a tiled dim with a partially-covering group so tiles cannot be represented.

Common situations: Collapsing leading dimensions of a tiled tensor where the tiling spans the collapse boundary with non-integer-multiple sizes.

Related errors

• Input/output tiling mismatch when attempting to collapse a s
• Reassociation {reassociation} is not compatible with tiling
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp

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