jax-ml/jax · error · ValueError

Cannot tile a transpose

Error message

Cannot tile a transpose ({permutation}). Tiling dims ({permutation[-tiling_len:]}) cannot contain non-tiled dims. All of them must be >= {tiling_offset}.

What it means

When transposing a tiled memref, the permuted tiling dimensions must only permute among themselves (i.e. each tiling dim index must be >= tiling_offset). Otherwise tiled and untiled dimensions would interleave, which the lowering cannot express.

Solutions

1. Restructure the permutation so the tiled (trailing) dimensions permute only among themselves
2. Apply the transpose before tiling, then re-tile the result
3. Choose tile shapes such that the desired permutation stays within the tiled block

Example fix

// before
# tiled dims are the trailing 2
out = t.memref.transpose(tiled, permutation=[2, 0, 1])  # mixes tile/untiled
// after
out = t.memref.transpose(tiled, permutation=[0, 2, 1])  # tile dims permuted among themselves

off = len(permutation) - tiling_len
assert all(d >= off for d in permutation[-tiling_len:]), 'tile dims must permute among themselves'

When it happens

Trigger: A permutation whose last tiling_len entries (the tiled dims) contain any index < len(permutation) - tiling_len, e.g. transposing a 2D tile across the tile/element boundary.

Common situations: Transposing a tiled smem buffer with a permutation that mixes tile dims with regular dims, e.g. permutation like [2,0,1] on shape with trailing tiling of 2.

Related errors

• Invalid in/out transforms. In transform
• Transpose cannot be moved before a tiling transform when it…
• Can not tile strides when tiled dimensions have been…
• 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/5bddfbeae15ab0aa. Report an issue: GitHub.

Appendix: source

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

    # 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}"
      )
    new_permutation = permutation + [
        x + tiling_len for x in permutation[-tiling_len:]
    ]
    new_permutation = _permutation_to_affine_map_attr(new_permutation)

  new_transpose_op = memref.TransposeOp(

View on GitHub (pinned to 1e1c6a8fc0)