jax-ml/jax · error · ValueError

Offsets {offsets} have lower rank than tiling {tiling}

Error message

Offsets {offsets} have lower rank than tiling {tiling}

What it means

tile_offset applies a tiling to the trailing dimensions of an offsets tuple; the offsets rank must be >= the tiling rank or tiling is undefined.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:999

  gmem_transforms: list[lc.MemRefTransform] = []
  for transform in attr:
    if isinstance(transform, mgpu.TileTransformAttr):
      tile_transform = lc.TileTransform(tuple(transform.tiling))
      gmem_transforms.append(tile_transform)
    elif not isinstance(transform, mgpu.SwizzleTransformAttr):
      raise NotImplementedError(f"Unsupported transform: {transform}")
  return tuple(gmem_transforms)


def tile_offset(
    offsets: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing offsets in `offsets` according to `tiling`.

  Raises if the offsets are not aligned with the start of a tile.
  """
  if len(offsets) < len(tiling):
    raise ValueError(f"Offsets {offsets} have lower rank than tiling {tiling}")
  untiled_offsets, tiled_offsets = (
      offsets[: -len(tiling)],
      offsets[-len(tiling) :],
  )
  for i, t in zip(tiled_offsets, tiling, strict=True):
    if i % t != 0:
      raise ValueError(f"Offset {i} is not divisible by tile size {t}")
  return (
      *untiled_offsets,
      *[i // t for i, t in zip(tiled_offsets, tiling, strict=True)],
      *[0] * len(tiling),
  )


def tile_strides(
    strides: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing strides in `strides` according to `tiling`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make len(offsets) >= len(tiling); give a full-rank offset tuple
2. Reduce the tiling rank to match the memref rank

Example fix

// before
tile_offset((4,), (8, 4))
// after
tile_offset((0, 4), (8, 4))

assert len(offsets) >= len(tiling), 'offsets rank must cover tiling rank'

When it happens

Trigger: Computing tile_offset(offsets, tiling) with fewer offsets than tiling dims (e.g. offsets=(4,) with tiling=(8,4)) via transform_type on a mis-shaped layout.

Common situations: Constructing tiled memref transforms where the tiling rank exceeds the memref rank, or passing a scalar offset for a 2D tile.

Related errors

• Strides {strides} have lower rank than tiling {tiling}
• Offset {i} is not divisible by tile size {t}
• Can not tile strides when tiled dimensions have been transpo
• Stride {s} is not divisible by {d} (tile size = {t}). Stride
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp

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