jax-ml/jax · error · ValueError

Reassociation {reassociation} is not compatible with tiling

Error message

Reassociation {reassociation} is not compatible with tiling {t_in.tiling}, as it causes tiled and untiled dimensions to be collapsed together

What it means

Defensive check in collapse_shape lowering: the reassociation must not cause tiled and untiled dimensions to be collapsed into the same output dimension. Walking reassociation groups from the end, the tile budget must never go negative.

Source

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

        f"Swizzle mismatch. In transforms swizzle: {in_swizzle}, out transforms"
        f" swizzle {out_swizzle}."
    )
  _check_collapse_shape(op, in_transforms, out_transforms)
  reassociation = [
      [ir.IntegerAttr(i).value for i in ir.ArrayAttr(dims)]
      for dims in op.reassociation
  ]
  new_reassociation = reassociation.copy()
  if in_transforms:
    [t_in] = in_transforms
    assert isinstance(t_in, lc.TileTransform)
    tiling_rank = to_process =  len(t_in.tiling)
    for index_from_end, dims in enumerate(reassociation[::-1]):
      to_process -= len(dims)
      if to_process < 0:
        # This should be caught by `_check_collapse_shape` today, but we check
        # it here as well in case `cs.CollapseShape` ever changes to allow this.
        raise ValueError(
            f"Reassociation {reassociation} is not compatible with tiling "
            f"{t_in.tiling}, as it causes tiled and untiled dimensions to "
            "be collapsed together"
        )
      if to_process == 0:
        for t_dims in reassociation[-index_from_end - 1:]:
          new_reassociation.append([dim + tiling_rank for dim in t_dims])
        break
    assert to_process == 0

  result = memref.collapse_shape(
      transform_type(op.result.type, out_transforms),
      unwrap_transformed_memref(op.src, in_transforms_attr),
      new_reassociation,
  )
  return [wrap_transformed_memref(result, op.result.type, out_transforms_attr)]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Regroup the reassociation so each output dim comes entirely from tiled or entirely from untiled source dims
2. Align tile shape to group boundaries
3. Untile before collapsing

to_process = tiling_rank
for dims in reassociation[::-1]:
    to_process -= len(dims)
    assert to_process >= 0, 'tiled and untiled dims collapsed together'

When it happens

Trigger: A reassociation group that mixes trailing tiled dimensions with non-tiled dimensions, e.g. group [1,2] where only dim 2 is tiled.

Common situations: Normally pre-empted by _check_collapse_shape, but reachable if the symbolic-collapse semantics change; user-visible when hand-crafting reassociations on tiled buffers.

Related errors

• Input tiling {t_in.tiling} is not compatible with {op}
• Input/output tiling mismatch when attempting to collapse a s
• 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/4f1427d8f78ddcc1. Report an issue: GitHub.