jax-ml/jax · error · ValueError

Tiles must have a decreasing rank

Error message

Tiles must have a decreasing rank

What it means

Tiling is a list of tile tuples where each successive tile must apply to no more dims than the previous one (non-increasing rank). This validation in Tiling.__post_init__ rejects configurations like [(2,), (4, 4)] after [(4, 4), (2,)].

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:79

  to the rank of the tile) is unfolded into two dimensions: first equal to the
  ratio of the dimension size and the tile size, and second equal to the tile
  size. Then, all newly unfolded minor dimensions are transposed to appear at
  the end.

  This expression describes multi-level tiling, by applying each element of
  `tiles` in sequence to the array.

  See https://openxla.org/xla/tiled_layout for a more detailed explanation.
  """
  tiles: tuple[tuple[int, ...], ...]

  def __post_init__(self):
    if not self.tiles:
      return
    last_tile_rank = len(self.tiles[0])
    for tile in self.tiles:
      if len(tile) > last_tile_rank:
        raise ValueError("Tiles must have a decreasing rank")
      if not tile:
        raise ValueError("Tiles must not be empty")
      if any(d <= 0 for d in tile):
        raise ValueError(f"Tile shape must only have positive sizes, got: {self.tiles}")
      last_tile_rank = len(tile)

  def __str__(self):
    return f"Tiling({''.join(map(str, self.tiles))})"

  def tile_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]:
    """Computes the shape of an array after tiling."""
    orig_shape = shape
    def fail():
      raise ValueError(f"Tiling {self.tiles} does not apply to shape {orig_shape}")
    for tile in self.tiles:
      if len(tile) > len(shape):
        fail()
      untiled_dims, tiled_dims = shape[:-len(tile)], shape[-len(tile):]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Order tiles outermost-first with non-increasing rank, e.g. [(8, 8), (4,), (2,)]
2. Build tilings via layout constructors (e.g. fa.TiledLayout) rather than literal tuples
3. Remember the first tile is the outermost covering the most dims

Example fix

// before
Tiling([(2,), (8, 8)])
// after
Tiling([(8, 8), (2,)])

ranks = [len(t) for t in tiles]
assert all(ranks[i] >= ranks[i+1] for i in range(len(ranks)-1)), 'tile ranks must be non-increasing'

def is_valid_tiling_order(tiles) -> bool:
    return all(len(tiles[i]) >= len(tiles[i+1]) for i in range(len(tiles)-1))

When it happens

Trigger: Constructing fragmented_array.Tiling with tiles whose ranks increase, e.g. Tiling([(2, 2), (4,) , (8, 8)]) — the inner (nested) tiles must have rank <= the outer ones.

Common situations: Hand-writing nested tilings for register layouts; converting a layout spec where tile order got reversed.

Related errors

• Tiles must not be empty
• Tile shape must only have positive sizes, got: {self.tiles}
• Shape {shape} and strides {strides} must have the same lengt
• Tiling {self.tiles} does not apply to shape {orig_shape}
• shape {orig_shape} is not a valid result of applying tiling

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