jax-ml/jax · error · ValueError

Can not tile strides when tiled dimensions have been transpo

Error message

Can not tile strides when tiled dimensions have been transposed with untiled dimensions. Strides: {strides}, tiling: {tiling}

What it means

Tiling assumes the trailing (tiled) dimensions carry the smallest strides; if the tiled dims have been transposed with untiled dims, the memory layout cannot be expressed as a simple tiling.

Source

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

def tile_strides(
    strides: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing strides in `strides` according to `tiling`.

  The `len(tiling)` trailing strides in `strides` must be the `len(tiling)`
  smallest strides in `strides`. The same property holds in the result, i.e.,
  given two tiles with indices i and j (i < j) with strides tiled according to
  this function, then all the elements in tile i are physically ordered before
  all the elements in tile j.

  E.g., tile_strides((2048, 32, 1), (8, 4)) = (2048, 256, 32, 4, 1)
  """
  if len(strides) < len(tiling):
    raise ValueError(f"Strides {strides} have lower rank than tiling {tiling}")
  ordered_strides = sorted(strides, reverse=True)
  if set(ordered_strides[-len(tiling):]) != set(strides[-len(tiling):]):
    raise ValueError(
        "Can not tile strides when tiled dimensions have been transposed with "
        f"untiled dimensions. Strides: {strides}, tiling: {tiling}"
    )
  untiled_strides, tiled_strides = strides[:-len(tiling)], strides[-len(tiling):]

  # Zip the strides and tiling together, in order to sort them together. This
  # allows handling cases where multiple tiling dimensions have the same stride,
  # which can occur with size-1 dimensions.
  tiled_strides_and_tiling: list[tuple[int, int]] = list(
      zip(tiled_strides, tiling, strict=True))
  tiled_ordered_strides_and_tiling = sorted(
      tiled_strides_and_tiling, reverse=True)

  to_ordered = lambda i: tiled_ordered_strides_and_tiling.index(tiled_strides_and_tiling[i])
  from_ordered = lambda i: tiled_strides_and_tiling.index(tiled_ordered_strides_and_tiling[i])

  ordered_tiling = [tiling[from_ordered(i)] for i in range(len(tiling))]
  ordered_tiled_strides = [tiled_strides[from_ordered(i)] for i in range(len(tiling))]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose the data in memory first so tiled dims are innermost with smallest strides
2. Apply the tile transform before the transpose
3. Use tile sizes of 1 on the transposed dims to effectively skip tiling them

Example fix

// before
tile_strides((1, 2048, 32), (8, 4))  # tiled dims transposed
// after
mem = transpose_to_row_major(mem)
tile_strides((2048, 32, 1), (8, 4))

ordered = sorted(strides, reverse=True)
assert set(ordered[-len(tiling):]) == set(strides[-len(tiling):]), 'tiled dims must be innermost contiguous'

When it happens

Trigger: tile_strides where the set of the len(tiling) smallest sorted strides differs from the set of the trailing strides, i.e. tiled dims are not the innermost-contiguous ones.

Common situations: Applying a tile transform to a transposed or column-major memref where inner strides don't correspond to the tiled dimensions.

Related errors

• Strides {strides} have lower rank than tiling {tiling}
• Stride {s} is not divisible by {d} (tile size = {t}). Stride
• `strides` must contain only 1s.
• Offsets {offsets} have lower rank than tiling {tiling}
• Offset {i} is not divisible by tile size {t}

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