jax-ml/jax · error · ValueError

Invalid dimension {dim} for tiling {self}

Error message

Invalid dimension {dim} for tiling {self}

What it means

Tiling.tile_dimension(dim) requires 0 <= dim < rank of the outermost tile (len(self.tiles[0])); the dim indexes into the tiled (outermost) dims, so out-of-range dims are invalid.

Source

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

      tiled_dims = shape[-len(canonical_tile):]
      if tiled_dims == canonical_tile:
        continue
      shape = canonical_tile
      new_tiling.append(canonical_tile)
    return Tiling(tuple(new_tiling))

  def tile_strides(self, strides: tuple[int, ...]) -> tuple[int, ...]:
    """Computes the strides of an array after tiling."""
    for tile in self.tiles:
      untiled, tiled = strides[:-len(tile)], strides[-len(tile):]
      strides = (*untiled, *(s * t for s, t in zip(tiled, tile)), *tiled)
    return strides

  def tile_dimension(self, dim: int) -> tuple[bool, ...]:
    """Result is True whenever the tiled dim originated from the given input dim."""
    tiling_rank = len(self.tiles[0])
    if dim < 0 or dim >= tiling_rank:
      raise ValueError(f"Invalid dimension {dim} for tiling {self}")
    strides = [1] * tiling_rank
    strides[dim] = 0
    return tuple(s == 0 for s in self.tile_strides(tuple(strides)))

  def remove_dimension(self, dim: int) -> Tiling:
    """Returns a tiling with the given dimension removed."""
    tiling_rank = len(self.tiles[0])
    if dim < 0 or dim >= tiling_rank:
      raise ValueError(f"Invalid dimension {dim} for tiling {self}")
    dim_in_tile = dim
    tiles = []
    last_tile_rank = len(self.tiles[0])
    for t in self.tiles:
      assert last_tile_rank >= len(t)
      dim_in_tile -= last_tile_rank - len(t)
      last_tile_rank = len(t)
      if dim_in_tile >= 0:
        t = t[:dim_in_tile] + t[dim_in_tile + 1:]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp/validate dim against len(tiling.tiles[0]) before calling
2. Recompute dim indices after layout transformations like remove_dimension
3. Use non-negative indices only

Example fix

// before
mask = tiling.tile_dimension(dim)
// after
assert 0 <= dim < len(tiling.tiles[0]), 'dim not in outer tile'
mask = tiling.tile_dimension(dim)

assert 0 <= dim < len(tiling.tiles[0])

def valid_tiled_dim(tiling, dim) -> bool:
    return 0 <= dim < len(tiling.tiles[0])

When it happens

Trigger: Calling tile_dimension(d) where d >= len(tiling.tiles[0]) or d < 0 — e.g. indexing a logical array dim that is not part of the outer tiling, or negative indexing which this API doesn't support.

Common situations: Reduction/broadcast lowering code computing dims from a different layout's rank; after removing dims the caller reuses stale dim indices.

Related errors

• Tiles must have a decreasing rank
• Tiles must not be empty
• Tile shape must only have positive sizes, got: {self.tiles}
• 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/1773b864cd025ca4. Report an issue: GitHub.