jax-ml/jax · error · ValueError

Expected tiling to be at most rank of shape. Got tiling: {ti

Error message

Expected tiling to be at most rank of shape. Got tiling: {tiling} (rank: {len(tiling)}) and shape {shape} (rank: {len(shape)}).

What it means

tile_shape tiles the trailing dimensions of a shape with a given tiling; the tiling's rank (its length) must not exceed the shape's rank. If tiling has more dimensions than shape, the per-dimension zip would silently mis-align, so the helper raises immediately.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:1784

  def refine(
      self,
      *,
      chunk: ir.Value | None = None,
      num_chunks: int | None = None,
      chunk_size: int | None = None,
  ):
    return Partition1D(
        self.partition.target_block_shape[0],
        num_chunks=num_chunks,
        chunk_size=chunk_size,
        base_offset=self.get_base(chunk) if chunk is not None else None,
    )


def tile_shape(shape, tiling):
  if len(tiling) > len(shape):
    raise ValueError(
        "Expected tiling to be at most rank of shape. Got tiling:"
        f" {tiling} (rank: {len(tiling)}) and shape {shape} (rank:"
        f" {len(shape)})."
    )
  if not tiling:
    return shape
  tiling_rank = len(tiling)
  for s, t in zip(shape[-tiling_rank:], tiling):
    if s % t:
      raise ValueError("Non-divisible tiling:", shape, tiling)
  return (
      *shape[:-tiling_rank],
      *(s // t for s, t in zip(shape[-tiling_rank:], tiling)),
      *tiling,
  )


def warp_tree_reduce(value, op, group_size):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the tiling to at most the rank of shape (drop leading tiling dims or align them to trailing dims)
2. Increase the shape rank if the tensor was meant to be multi-dimensional
3. Check where the tiling was constructed (e.g., a block spec) and fix the rank mismatch at the source

Example fix

# before
tile_shape(shape=(4096,), tiling=(8, 128))
# after
tile_shape(shape=(4096,), tiling=(128,))

assert len(tiling) <= len(shape), f'tiling rank {len(tiling)} > shape rank {len(shape)}'
result = tile_shape(shape, tiling)

When it happens

Trigger: Calling tile_shape(shape=(256,), tiling=(8, 8)) or any call site where the tiling tuple/list is longer than the shape tuple. Common in kernels configured with a 2D warp tile applied to 1D buffers.

Common situations: Copy-pasting a 2D matmul tiling config into a 1D reduction/normalization kernel; changing tensor rank without updating the tiling constants; mismatch between layout rank and block tiling in a Mosaic kernel.

Related errors

• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w

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