jax-ml/jax · error · ValueError

Size mismatch for in/out transforms. In transforms: {in_tran

Error message

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

What it means

memref.transpose lowering requires the number of MemRefTransforms on input and output to be equal. Extra or missing transforms on either side cannot be reconciled by permutation alone.

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)

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

• Transpose cannot be moved before a tiling transform when it
• Swizzle mismatch. In transforms swizzle: {in_swizzle}, out t
• Expected the same number of in/out transforms, but got {in_t
• 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/08b2d4fbb7e09746. Report an issue: GitHub.