jax-ml/jax · error · ValueError

Tiling {self.tiles} does not apply to shape {orig_shape}

Error message

Tiling {self.tiles} does not apply to shape {orig_shape}

What it means

Tiling.tile_shape applies each tile to the trailing dims of the shape; it fails when a tile has more dims than remain, or the shape dims are not evenly divisible by the corresponding tile dims.

Source

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

      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):]
      if any(s % t != 0 for s, t in zip(tiled_dims, tile)):
        fail()
      shape = (*untiled_dims, *(d // t for d, t in zip(tiled_dims, tile)), *tile)
    return shape

  def untile_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]:
    """Computes the shape of an array before tiling from its tiled shape."""
    orig_shape = shape
    def fail():
      raise ValueError(
          f"shape {orig_shape} is not a valid result of applying tiling {self}."
      )
    for tile in reversed(self.tiles):
      if len(tile) > len(shape):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or slice the array so trailing dims are multiples of the tile dims
2. Use a tiling whose tiles match the array's shape (smaller final tile)
3. Check shape lengths: the outermost (first) tile must have rank <= array rank

Example fix

// before
Tiling([(4, 4)]).tile_shape((6, 8))  # 6 % 4 != 0
// after
Tiling([(4, 4), (2,)]).tile_shape((6, 8))  # 6 = 2*4... use tile matching shape

def applies(t, shape):
    s = shape
    for tile in t.tiles:
        if len(tile) > len(s): return False
        if any(x % y for x, y in zip(s[-len(tile):], tile)): return False
        s = s[:-len(tile)] + tuple(x//y for x, y in zip(s[-len(tile):], tile))
    return True
assert applies(tiling, shape)

When it happens

Trigger: Calling tile_shape((3, 8)) with Tiling([(4, 4)]) — 3 not divisible by 4; or tile_shape((4,)) with a rank-2 tile since len(tile) > len(shape).

Common situations: Applying a register layout designed for one tensor shape to a differently-shaped tensor; padding assumptions that don't hold.

Related errors

• shape {orig_shape} is not a valid result of applying tiling
• Tiles must have a decreasing rank
• Tiles must not be empty
• Tile shape must only have positive sizes, got: {self.tiles}
• Invalid dimension {dim} for tiling {self}

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