jax-ml/jax · error · NotImplementedError

Sum of sizes {n} must be equal to dimension {axis} of the op

Error message

Sum of sizes {n} must be equal to dimension {axis} of the operand shape {list(aval_in.shape)}.

What it means

Thrown by the Pallas fuser's split-operand rule when splitting one operand into multiple outputs: the sum of the requested split sizes along an axis must exactly equal that axis's length in the input's shape. It fires when sizes (e.g. from lax.split or a custom split usage rule) don't add up to aval_in.shape[axis].

Source

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


@register_pull_block_spec_rule(lax.split_p)
def _split_pull_rule(
    ctx: PullRuleContext,
    out_block_transforms: tuple[BlockIndexTransform, ...],
    *,
    sizes: Sequence[int],
    axis: int,
):
  aval_in = ctx.avals_in[0]
  assert isinstance(aval_in, core.ShapedArray)
  assert all(isinstance(aval, core.ShapedArray) for aval in ctx.avals_out)

  # turn numpy ints into ints
  sizes = [int(s) for s in sizes]
  n = sum(sizes)
  if n != aval_in.shape[axis]:
    raise NotImplementedError(
        f'Sum of sizes {n} must be equal to dimension {axis} of the operand '
        f'shape {list(aval_in.shape)}.'
    )
  valid_transforms = [
      bt for bt in out_block_transforms if bt is not no_block_index_transform
  ]
  if not valid_transforms:
    return [no_block_index_transform]

  block_transform = valid_transforms[0]

  new_block_shape = list(block_transform.block_shape)
  new_block_shape[axis] = pallas_core.Blocked(n)

  def new_block_index_transform(*idxs):
    idx = list(block_transform.block_index_transform(*idxs))
    idx[axis] = 0
    return tuple(idx)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Verify sum(sizes) == operand.shape[axis] before the split (e.g. use jnp.array_split semantics only when even division is guaranteed)
Use explicit size lists instead of a bare integer num-chunks when the axis length isn't divisible
Move the split outside the fused region and pass pre-split arrays into the kernel

Example fix

// before
parts = lax.split(x, [2])  # sizes=[2] but x.shape[0] == 5
// after
parts = lax.split(x, [2, 5])  # sizes sum (2+3) == x.shape[0]

Defensive patterns

Strategy: validation

Validate before calling

sizes = [int(s) for s in sizes]
assert sum(sizes) == x.shape[axis], f'{sum(sizes)} != {x.shape[axis]}'

Type guard

def split_sizes_valid(x, sizes, axis) -> bool:
    return sum(int(s) for s in sizes) == x.shape[axis]

Try / catch

try:
    parts = fused_split(x)
except NotImplementedError as e:
    if 'Sum of sizes' in str(e):
        parts = jnp.split(x, indices_or_sections=len(x.shape[axis] and sizes), axis=axis)
    else:
        raise

Prevention

Always derive split points from the actual shape: np.cumsum(sizes)[:-1]
Use jnp.array_split only when uneven splits are acceptable outside fusion
Assert divisibility before integer-chunk splits

When it happens

Trigger: Using lax.split / jnp.split-style operations inside a fused Pallas region where the split sizes don't sum to the operand's dimension length, or where a size is a numpy int that truncates (though these are coerced with int(s) first).

Common situations: Computing split sizes dynamically and off by one; splitting a dimension whose length isn't divisible by the number of chunks; passing sizes in the wrong order or against the wrong axis index.

Related errors

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