jax-ml/jax · error · ValueError

Input/output tiling mismatch when attempting to collapse a s

Error message

Input/output tiling mismatch when attempting to collapse a shape. Expected output tiling to be {expected_t_out.tiling} for input tiling {t_in.tiling}, but got {t_out.tiling}

What it means

For a collapsible tiled memref, the output tiling is fully determined by the input tiling and the reassociation. This error fires when the annotated output tiling differs from the computed expected tiling.

Source

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

        "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

  [in_transforms_attr] = inference_utils.in_transforms(op)
  [out_transforms_attr] = inference_utils.out_transforms(op)

  in_swizzle = swizzle_from_transforms_attr(in_transforms_attr)
  in_transforms = memref_transforms_from_transforms_attr(in_transforms_attr)
  out_swizzle = swizzle_from_transforms_attr(out_transforms_attr)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Let inference compute the output tiling (omit the manual out transform)
2. Or set the out tiling to the expected collapsed tiling reported in the message
3. Recompute: run cs.reduce_expression(cs.CollapseShape(...)) mentally/externally to get the right value

expected = cs.reduce_expression(cs.CollapseShape(cs.SMEMTransforms(t_in, None), tuple(shape), reassoc), {})
assert not isinstance(expected, cs.Unsatisfiable) and expected.tiling.tiling == t_out.tiling

When it happens

Trigger: memref.collapse_shape with a manually annotated out transform whose tiling != collapsed_tiling.tiling derived from t_in and the reassociation.

Common situations: Hand-writing the output transform instead of letting inference derive it, causing a mismatch with the canonical collapsed tiling.

Related errors

• Invalid in/out transforms. In transform: {in_transform}, out
• Input tiling {t_in.tiling} is not compatible with {op}
• Reassociation {reassociation} is not compatible with tiling
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is

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