jax-ml/jax · error · ValueError

Strides {strides} have lower rank than tiling {tiling}

Error message

Strides {strides} have lower rank than tiling {tiling}

What it means

tile_strides tiles the trailing strides of a memref; the strides tuple must have rank >= tiling rank, otherwise the tiling is undefined.

Source

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

      *[0] * len(tiling),
  )


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])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Match tiling rank to memref rank (len(strides))
  2. Pass a full-rank strides array derived from the memref layout

Example fix

// before
tile_strides((1,), (8, 4))
// after
tile_strides((4, 1), (8, 4))
Defensive patterns

Strategy: validation

Validate before calling

assert len(strides) >= len(tiling), 'strides rank must cover tiling rank'

Prevention

When it happens

Trigger: Computing tile_strides(strides, tiling) with len(strides) < len(tiling), typically from a memref layout whose rank is smaller than the tile dims.

Common situations: Applying a 2D tile transform to a 1D memref, or a strides/tiling rank mismatch when constructing transform_type.

Related errors


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