jax-ml/jax · error · ValueError

Swizzle mismatch. In transforms swizzle: {in_swizzle}, out t

Error message

Swizzle mismatch. In transforms swizzle: {in_swizzle}, out transforms swizzle {out_swizzle}.

What it means

The transpose lowering requires the swizzle component of the input and output transform annotations to be identical. Transpose permutes dimensions but cannot re-derive a different swizzle pattern, so mismatched swizzles raise this ValueError.

Source

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

  return ir.AffineMapAttr.get(ir.AffineMap.get_permutation(permutation))


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

  in_transforms_attr = inference_utils.in_transforms(op)[0]
  unwrapped_in_ref = unwrap_transformed_memref(op.in_, in_transforms_attr)
  in_swizzle = swizzle_from_transforms_attr(in_transforms_attr)
  in_transforms = memref_transforms_from_transforms_attr(in_transforms_attr)
  out_transforms_attr = inference_utils.out_transforms(op)[0]
  out_swizzle = swizzle_from_transforms_attr(out_transforms_attr)
  out_transforms = memref_transforms_from_transforms_attr(out_transforms_attr)

  if in_swizzle != out_swizzle:
    raise ValueError(
        f"Swizzle mismatch. In transforms swizzle: {in_swizzle}, out transforms"
        f" swizzle {out_swizzle}."
    )
  if len(out_transforms) != len(in_transforms):
    raise ValueError(
        f"Size mismatch for in/out transforms. In transforms: {in_transforms},"
        f" out transforms: {out_transforms}."
    )
  if not out_transforms:
    new_permutation = op.permutation
  else:
    permutation = [
        ir.AffineDimExpr(e).position
        for e in op.permutation.value.results
    ]
    # We expect to have the same transforms on in/out, up to permutation of the
    # out transforms.
    # For example, for 3D input:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep the same swizzle on both sides of the transpose; apply a separate op to change swizzle
2. Drop swizzle from both annotations if the layout doesn't need it
3. Check swizzle_from_transforms_attr on both attrs and align them before building the transpose

assert swizzle_from_transforms_attr(in_attr) == swizzle_from_transforms_attr(out_attr)

When it happens

Trigger: memref.transpose with in_transforms and out_transforms whose swizzle_from_transforms_attr values differ (e.g. input swizzle=SWIZZLE_32B, output swizzle=SWIZZLE_128B or None).

Common situations: Transposing a swizzled shared-memory tile while also trying to change its swizzle mode in one op; common when tuning shared-memory layouts for MMA pipelines.

Related errors

• Transpose cannot be moved before a tiling transform when it
• Can't transpose the swizzled dimension.
• Size mismatch for in/out transforms. In transforms: {in_tran
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is

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