jax-ml/jax · error · ValueError

Size mismatch for in/out transforms. In transforms

Error message

Size mismatch for in/out transforms. In transforms: {in_transforms}, out transforms: {out_transforms}.

What it means

Raised by the memref.TransposeOp lowering when the input and output transform lists have different lengths. In/out transforms must correspond one-to-one (up to permutation) for a transpose to be lowered.

Solutions

1. Set identical transform lists (same count) on input and output of the transpose
2. Use Mosaic's high-level transpose/permute helper so inference fills both sides consistently
3. If no transforms are intended, clear the annotations on both sides

assert len(in_transforms) == len(out_transforms), 'transform count mismatch'

When it happens

Trigger: Building a transpose where in_transforms has, say, one tile transform but out_transforms has zero or two (or vice versa).

Common situations: Transform inference producing a default empty transform on one side while the other is explicitly annotated; typically from manually constructing the op rather than using Mosaic's Python builder APIs.

Related errors

• Expected the same number of in/out transforms, but got
• Swizzle mismatch. In transforms swizzle
• Transpose cannot be moved before a tiling transform when it…
• Can't transpose a TMEM reference.
• Can't transpose the swizzled dimension.

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

Appendix: source

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

    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:
    #   permutation: (0, 2, 1)
    #   in_transforms: TilingTransform((32, 8))
    # We expect:
    #   out_transforms: TilingTransform((8, 32))
    # TODO(olechwierowicz): Support multiple transforms.

View on GitHub (pinned to 1e1c6a8fc0)