jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

When a transpose carries transforms, both the input and output transform must be TileTransforms (only tiling is supported, not swizzle or other transform kinds, and only a single one each).

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)

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
• for transpose support, subclass {type(self)} must implement
• Cannot permute last two dimensions with leading dimensions.
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.

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