jax-ml/jax · error · ValueError

shape {orig_shape} is not a valid result of applying tiling

Error message

shape {orig_shape} is not a valid result of applying tiling {self}.

What it means

untile_shape inverts tiling: a tiled shape has extra tiled dims inserted; this error fires when the given shape doesn't have room for the tile's dims (len mismatch) or the tiled-dim structure doesn't match the tiling.

Source

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

  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):
        fail()
      untiled_dims = shape[:-2 * len(tile)]
      tiled_dims = shape[-2 * len(tile):-len(tile)]
      tiling_dims = shape[-len(tile):]
      if tiling_dims != tile:
        fail()
      shape = (*untiled_dims, *(d * t for d, t in zip(tiled_dims, tile)))
    return shape

  def canonicalize(self) -> Tiling:
    """Returns a canonicalized version of the tiling.

    We define a tiling to be canonical if, at each step (except the first one,
    which defines the base tile shape):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Only call untile_shape with shapes produced by tile_shape under the same Tiling
2. Verify the tiling instance matches the one used to tile originally
3. Reconstruct via the layout object rather than manual shape math

tiled = tiling.tile_shape(orig_shape)
# only untile shapes equal to `tiled` for the same tiling

When it happens

Trigger: Calling untile_shape with a shape that wasn't produced by tile_shape for this tiling — e.g. wrong number of dims for reversed tiles or non-matching dim pairs.

Common situations: Round-tripping shapes through layout transformations where an intermediate op changed the shape; mixing layouts between tensors of different shapes.

Related errors

• Tiling {self.tiles} does not apply to shape {orig_shape}
• 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/dae10c4390591a17. Report an issue: GitHub.