jax-ml/jax · error · ValueError

Offsets have lower rank than tiling

Error message

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

What it means

tile_offset tiles trailing offsets, which requires at least as many offset entries as there are tiling levels so each level has an offset to split. It fires when len(offsets) < len(tiling), leaving some tile levels without a corresponding offset.

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 have lower rank than tiling
• Can not tile strides when tiled dimensions have been…
• Offset is not divisible by tile size
• Stride is not divisible by (tile size = ). Strides: , tiling
• A scale address calculation for multiple M tiles

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

Appendix: 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)