jax-ml/jax · error · ValueError

Cannot tile a transpose ({permutation}). Tiling dims ({permu

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.

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)

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

• Transpose cannot be moved before a tiling transform when it
• Invalid in/out transforms. In transform: {in_transform}, out
• Cannot permute last two dimensions with leading dimensions.
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp

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