jax-ml/jax · error · ValueError

Expected the same number of in/out transforms, but got {in_t

Error message

Expected the same number of in/out transforms, but got {in_transforms=} and {out_transforms=}

What it means

memref.collapse_shape lowering requires exactly the same number of input and output transforms. A mismatch means the layout annotations are inconsistent and the collapse semantics are undefined.

Source

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

      output_shape=op.output_shape,
      static_output_shape=out_transformed_ty.shape,
  )

  wrapped_ref = wrap_transformed_memref(
      new_expand_shape_op.result, op.result.type, out_transforms
  )
  return [wrapped_ref]


# TODO(bchetioui): find a way to consolidate the logic that is shared logic with
# layout inference. It is not entirely clear what the best approach is.
def _check_collapse_shape(
    op: memref.CollapseShapeOp,
    in_transforms: Sequence[lc.MemRefTransform],
    out_transforms: Sequence[lc.MemRefTransform],
):
  if len(in_transforms) != len(out_transforms):
    raise ValueError(
        "Expected the same number of in/out transforms, but got "
        f"{in_transforms=} and {out_transforms=}"
    )
  if not in_transforms:
    return
  t_in, *in_transforms = in_transforms
  t_out, *_ = out_transforms
  if (in_transforms or
      not isinstance(t_in, lc.TileTransform) or
      not isinstance(t_out, lc.TileTransform)):
    raise NotImplementedError(
        "Only a single tiling transform is supported when collapsing a shape, "
        f"but got {in_transforms=} and {out_transforms=}"
    )
  src_ty = ir.MemRefType(op.src.type)
  strides, _ = src_ty.get_strides_and_offset()
  if strides != utils.get_contiguous_strides(src_ty.shape):
    raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Mirror the transform count: same number of transforms (typically one TileTransform) on both sides
2. Omit manual annotations and rely on inference for the output
3. For untransformed memrefs ensure both annotations are empty

assert len(in_transforms) == len(out_transforms)

When it happens

Trigger: memref.collapse_shape where len(in_transforms) != len(out_transforms), e.g. input annotated with one tiling and output with none.

Common situations: Collapsing a tiled smem buffer while manually constructing the result type without carrying the transform annotation over.

Related errors

• Size mismatch for in/out transforms. In transforms: {in_tran
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unsupported transform: {type(transform)}

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