jax-ml/jax · error · NotImplementedError

Block size must be a multiple of the input size. Got block {

Error message

Block size must be a multiple of the input size. Got block {block_shape=} but input {x.shape}.

What it means

The tile rule additionally requires each output block dimension to be an integer multiple of the corresponding input dimension. If out_dim % in_dim != 0 the repeat count is fractional, which cannot be expressed, so the fuser raises NotImplementedError.

Source

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

  del eval_ctx
  return lax.split(x, sizes=sizes, axis=axis)


@register_eval_rule(lax.tile_p)
def _tile_eval_rule(
    eval_ctx: KernelEvalContext, x, reps: tuple[int, ...]
):
  block_spec = eval_ctx.out_block_specs[0]
  block_shape = tuple(d for d in block_spec.block_shape
                      if not isinstance(d, pallas_core.Squeezed))
  if not all(isinstance(dim, int) for dim in block_shape):
    raise NotImplementedError(
        'tile with non-int block dimensions not supported yet'
    )
  if not all(
      out_dim % in_dim == 0 for out_dim, in_dim in zip(block_shape, x.shape)
  ):
    raise NotImplementedError(
        'Block size must be a multiple of the input size. '
        f'Got block {block_shape=} but input {x.shape}.'
    )
  reps_in_block = [
      out_dim // in_dim if out_dim >= in_dim else 1
      for out_dim, in_dim in zip(block_shape, x.shape)
  ]
  return lax.tile(x, reps_in_block)


@register_pull_block_spec_rule(lax.tile_p)
def _tile_pull_rule(
    ctx: PullRuleContext,
    block_transform: BlockIndexTransform,
    *,
    reps: tuple[int, ...],
):
  del reps

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the output block shape an exact multiple of the input shape on every axis (adjust block size or pad the input first)
2. Pad the input (e.g. to the next multiple) before tiling
3. Hoist the tile out of the fused region

Example fix

// before
y = fuse(jnp.tile)(x, (8, 1))  # x.shape[0] == 3, 8 % 3 != 0
// after
x_padded = jnp.pad(x, ((0, 1), (0, 0)))  # now length 4
y = fuse(jnp.tile)(x_padded, (8, 1))  # 8 % 4 == 0

assert all(b % i == 0 for b, i in zip(block_shape, x.shape) if i > 0), 'block must be multiple of input on every axis'

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

try:
    y = fused_tile(x, reps)
except NotImplementedError as e:
    if 'multiple of the input size' in str(e):
        y = jnp.tile(x, reps)  # or pad input to a multiple first
    else:
        raise

When it happens

Trigger: jnp.tile / lax.tile inside a fused Pallas region where the output block shape is not a whole multiple of the input shape along at least one axis, e.g. tiling a length-3 axis into a block of 8.

Common situations: Mismatched tiling factors between BlockSpec and the tile call; padding targets whose size isn't divisible by the input extent; combining squeeze/reshape with tile so shapes drift.

Related errors

• Every block dimension must be either a multiple or factor of
• Block shape for {origin} (= {block_shape}) must have the sam
• Stacking only supported when the block size along the stack
• Sum of sizes {n} must be equal to dimension {axis} of the op
• tile with non-int block dimensions not supported yet

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