jax-ml/jax · error · ValueError

Transpose cannot be moved before a tiling transform when it

Error message

Transpose cannot be moved before a tiling transform when it changes the set of tiled dimensions. (permutation: {perm}, tiling: {self.tiling})

What it means

When a transpose transform is commuted past a TilingTransform, the transpose must permute only the non-tiled leading dimensions — it cannot change which dimensions are tiled. If the permutation, offset by the untiled rank, does not map onto the tiled suffix of dimensions, ValueError is raised because the tiling could not be consistently re-derived.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:786

      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise TypeError(f"Cannot transform type: {x}")

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    return TilingTransform(self.tiling)

  def commute_transpose(
      self, _: jax_core.AbstractValue,
      transpose: state_types.TransposeTransform,
  ) -> tuple[state_types.TransposeTransform, UntilingTransform]:
    # The transpose in question is applied to the untiled ref so we
    # need to translate it by duplicating and offsetting the last part.
    perm = transpose.permutation
    off = len(perm)
    new_suffix = [i + off for i in perm[-len(self.tiling) :]]
    if set(new_suffix) != set(range(off, off + len(self.tiling))):
      raise ValueError(
          "Transpose cannot be moved before a tiling transform when it changes"
          f" the set of tiled dimensions. (permutation: {perm}, tiling:"
          f" {self.tiling})"
      )

    new_tiling = tuple(self.tiling[i - off] for i in new_suffix)
    new_transpose = state_types.TransposeTransform((*perm, *new_suffix))
    return new_transpose, dataclasses.replace(self, tiling=new_tiling)

  def commute_ndindexer(
      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, UntilingTransform]:
    del aval
    idxs = indexer.indices
    indexer_shape = indexer.shape
    untiled_idxs = idxs[: -len(self.tiling)]
    tiled_idxs = idxs[-len(self.tiling) :]
    idxs_after_tiling: list[indexing.Slice] = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose before tiling, i.e. apply the transpose to the untiled ref so the tiled dimensions stay last
2. Restructure the kernel to index the transposed view explicitly instead of commuting transforms
3. Choose a tiling whose tiled dimensions are invariant under your permutation (permute only leading dims)

Example fix

# before
ref_t = tiling_transform.apply(ref)
ref_tt = transpose_transform.apply(ref_t)  # ValueError if tiled dims move

# after
ref_t = transpose_transform.apply(ref)   # transpose first, on untiled ref
ref_tt = tiling_transform.apply(ref_t)    # tiling applied after

# ensure the transpose only permutes leading (untiled) dims
n_tiled = len(tiling)
assert set(perm[-n_tiled:]) == set(
    range(len(perm) - n_tiled, len(perm)
)), 'transpose moves tiled dimensions'

def transpose_preserves_tiled_dims(perm, n_tiled) -> bool:
    return sorted(perm[-n_tiled:]) == list(
        range(len(perm) - n_tiled, len(perm))
    )

null

When it happens

Trigger: Applying a transpose on a tiled ref that moves a tiled dimension into a leading position (or vice versa) and then triggering transform commutation during lowering, e.g. transposing a 2D tiled block spec before passing to a kernel.

Common situations: Transposing blocked matrices in Pallas kernels where the tile layout is fixed (e.g. swapping rows/cols of a (128, N)-tiled TMEM accumulator); chained transforms (tiling then transpose) produced by automatic transform compositions in newer JAX versions.

Related errors

• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Unsupported transform: {type(transform)}

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