jax-ml/jax · error · ValueError

Expected slice start ({start}) and slice size ({size}) to be

Error message

Expected slice start ({start}) and slice size ({size}) to be divisible by the tile size ({tile})

What it means

For a slice through a tiled dimension, TilingTransform.commute_ndindexer requires the slice's static start and size to be divisible by the tile size, so the slice can be rewritten exactly into tiled coordinates. Otherwise ValueError is raised with the offending start/size and tile.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:824

        indexer_shape[: -len(self.tiling)],
        indexer_shape[-len(self.tiling) :],
    )
    for idx, tile, dim in zip(tiled_idxs, self.tiling, untiled_shape):
      match idx:
        case slice() | indexing.Slice():
          if isinstance(idx, slice):
            ds = indexing.Slice.from_slice(idx, dim)
          else:
            ds = idx
          if ds.stride is not None and ds.stride != 1:
            raise NotImplementedError(
                f"Strided slices unsupported. Got stride: {ds.stride}"
            )
          start, size = ds.start, ds.size
          if (
              start is not None and isinstance(start, int) and start % tile
          ) or (size is not None and isinstance(size, int) and size % tile):
            raise ValueError(
                f"Expected slice start ({start}) and slice size ({size})"
                f" to be divisible by the tile size ({tile})"
            )
          def _maybe_cdiv_with_cast(x, y):
            if x is None:
              return None
            if isinstance(x, jax.Array):
              # If x is an int32, we need to make sure y is an int32 to avoid
              # a dtype mismatch.
              y = jnp.array(y, x.dtype)
            return pallas_utils.cdiv(x, y)
          new_start = _maybe_cdiv_with_cast(start, tile)
          new_size = _maybe_cdiv_with_cast(size, tile)
          idxs_after_tiling.append(indexing.Slice(new_start, new_size))
        case _:
          raise TypeError(f"Unsupported index type: {type(idx)}")
    assert all(a % b == 0 for a, b in zip(untiled_shape, self.tiling))
    tiled_shape = [

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or round slice start and size up/down to multiples of the tile size
2. Choose a tile size that divides your access pattern (e.g. tile=16 for 16-aligned blocks)
3. Use integer indexing on the tiled dim with computed tile indices instead of raw slices

Example fix

# before
x = ref[16:48]  # start 16, size 32; tile=32 -> start not divisible

# after
TILE = 32
start = (16 // TILE) * TILE        # align to tile boundary
x = ref[start:start + 2 * TILE]    # aligned start and tile-multiple size

start, size, tile = 16, 32, 32
aligned = (start % tile == 0) and (size is None or size % tile == 0)

def slice_aligned_to_tile(start, size, tile) -> bool:
    return (start is None or start % tile == 0) and (
        size is None or size % tile == 0
    )

null

When it happens

Trigger: Slicing a tiled ref with offsets/sizes not aligned to the tile, e.g. tile size 32 and slice `ref[16:48]` (start 16 % 32 != 0) or size 48 not a multiple of 32.

Common situations: Block kernels with non-aligned offsets; dynamic block sizes producing unaligned slices; migrating hand-written SMEM code to tiled block specs where alignment was implicit.

Related errors

• Strided slices unsupported. Got stride: {ds.stride}
• The product of the major dimensions must be a multiple of {i
• The Pallas TPU lowering currently requires that the last two
• Get only supports slices with stride 1, got {strides}
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp

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