jax-ml/jax · error · ValueError

Can not tile strides when tiled dimensions have been…

Error message

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

What it means

Raised by tile_strides when the tiled (trailing) dimensions have been transposed relative to the untiled dimensions: tiling requires the tiled strides to be the smallest, trailing ones. Fires when that ordering invariant (checked via sorted stride comparison) is violated.

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

• Stride is not divisible by (tile size = ). Strides: , tiling
• Strides have lower rank than tiling
• Offset is not divisible by tile size
• Offsets have lower rank than tiling
• `strides` must contain only 1s.

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

Appendix: 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)