jax-ml/jax · error · NotImplementedError

Every block dimension must be either a multiple or factor of

Error message

Every block dimension must be either a multiple or factor of input. Got block {block_shape} for input {aval_in.shape}

What it means

Companion check to 1886: for tiling, every block dimension's effective size (via pallas_core.get_block_size) must be either an exact multiple or an exact factor of the corresponding input dimension. Anything else (e.g. block 6 vs input 4) can't be tiled whole-number-of-times, so the rule fails.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:2143

    reps: tuple[int, ...],
):
  del reps
  block_shape = block_transform.block_shape
  aval_in = ctx.avals_in[0]
  assert isinstance(aval_in, core.ShapedArray)
  assert len(block_shape) == len(aval_in.shape)
  if not all(isinstance(dim, (int, pallas_core.Squeezed))
             for dim in block_shape):
    raise NotImplementedError(
        'tile with non-int block dimensions not supported yet'
    )

  if not all(
      (pallas_core.get_block_size(block_dim) % in_dim == 0) or
      (in_dim % pallas_core.get_block_size(block_dim) == 0)
      for block_dim, in_dim in zip(block_shape, aval_in.shape)
  ):
    raise NotImplementedError(
        'Every block dimension must be either a multiple or factor of input. '
        f'Got block {block_shape} for input {aval_in.shape}'
    )

  new_shape = tuple(
      block_dim if isinstance(block_dim, pallas_core.Squeezed)
      else min(block_dim, in_dim)
      for block_dim, in_dim in zip(block_shape, aval_in.shape)
  )

  def new_block_index_transform(*idxs):
    original_idxs = block_transform.block_index_transform(*idxs)
    return tuple(
        0 if pallas_core.get_block_size(block_dim) >= in_dim
        else orig_idx % (in_dim // pallas_core.get_block_size(block_dim))
        for orig_idx, block_dim, in_dim in zip(
            original_idxs, block_shape, aval_in.shape
        )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Choose block sizes that are exact multiples of the input dims on tiled axes (or exact factors)
Reshape/pad the input so each axis divides or is divided by the block size
Skip tiling on that axis by using a block dim equal to the input dim

Example fix

// before
block_shape = (6, 128)  # input shape (4, 128): 6 % 4 != 0 and 4 % 6 != 0
// after
block_shape = (8, 128)  # 8 % 4 == 0

Defensive patterns

Strategy: validation

Validate before calling

from jax._src.pallas import core as pc
assert all((pc.get_block_size(b) % i == 0) or (i % pc.get_block_size(b) == 0)
           for b, i in zip(block_shape, x.shape)), 'block dims must be multiple/factor of input'

Type guard

def block_divides_or_divisible(block_shape, in_shape) -> bool:
    return all((b % i == 0) or (i % b == 0) for b, i in zip(block_shape, in_shape))

Try / catch

try:
    out = fused_tile_fn(x)
except NotImplementedError as e:
    if 'multiple or factor' in str(e):
        out = unfused_tile(x)
    else:
        raise

Prevention

Pick hardware tile sizes that are powers of 2 dividing/ multiples of operand extents
Reshape operands so extents divide block sizes
Validate block/operand divisibility in a pre-launch assert

When it happens

Trigger: Tile/broadcast usage inside a fused Pallas kernel where for some axis neither block_size % in_dim == 0 nor in_dim % block_size == 0 holds (Squeezed dims contribute size 1, which always divides).

Common situations: Block sizes chosen for hardware tiling (e.g. 128-wide tiles) applied to operands whose extents share no divisor relationship; partial reshapes leaving odd extents; changing block sizes without re-checking operand shapes.

Related errors

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