jax-ml/jax · error · NotImplementedError

Unsupported in/out transforms. In transform

Error message

Unsupported in/out transforms. In transform: {in_transform}, out transform: {transform}

What it means

Raised by the memref.TransposeOp lowering when the paired in/out transforms are not both TileTransforms (or otherwise not in the supported correspondence). Only matched tiling transforms (up to permutation) are supported for transposed memrefs.

Solutions

1. Reduce transforms on both sides to a single TileTransform each
2. Move swizzle/other transforms outside the transpose
3. If you need multiple transforms around a transpose, materialize them as explicit reshape/layout ops instead

def both_single_tile_transforms(in_ts, out_ts) -> bool:
    return (len(in_ts) == 1 and len(out_ts) == 1
            and isinstance(in_ts[0], lc.TileTransform)
            and isinstance(out_ts[0], lc.TileTransform))

When it happens

Trigger: memref.transpose with non-empty transforms where either in_transform or out_transform is not an instance of lc.TileTransform (e.g. a swizzle-only or slice transform).

Common situations: Stacking a swizzle annotation onto a transposed tiled tensor; hitting the documented TODO(olechwierowicz) limitation of single-tile-transform support.

Related errors

• Only a single tiling transform is supported when collapsing…
• axes argument to transpose()
• Can't transpose a TMEM reference.
• Can't transpose the swizzled dimension.
• Cannot convert from to

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

Appendix: source

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

  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.
    [transform] = out_transforms
    [in_transform] = in_transforms
    if not isinstance(transform, lc.TileTransform) or not isinstance(
        in_transform, lc.TileTransform
    ):
      raise NotImplementedError(
          f"Unsupported in/out transforms. In transform: {in_transform}, out"
          f" transform: {transform}"
      )
    tiling_len = len(in_transform.tiling)
    tiling_offset = len(permutation) - tiling_len
    if any(dim < tiling_offset for dim in permutation[-tiling_len :]):
      raise ValueError(
          f"Cannot tile a transpose ({permutation}). Tiling dims"
          f" ({permutation[-tiling_len:]}) cannot contain non-tiled dims."
          f" All of them must be >= {tiling_offset}."
      )
    dims = [-1] * len(permutation)
    dims = dims[:-tiling_len] + list(in_transform.tiling)
    permuted_dims = tuple(dims[permutation[i]] for i in range(len(dims)))
    if permuted_dims[-tiling_len:] != transform.tiling:
      raise ValueError(
          f"Invalid in/out transforms. In transform: {in_transform}, out"
          f" transform: {transform}"

View on GitHub (pinned to 1e1c6a8fc0)