{"record":{"id":"98c0ec67361eedcf","repo":"jax-ml/jax","slug":"sum-of-sizes-n-must-be-equal-to-dimension-axis","errorCode":null,"errorMessage":"Sum of sizes {n} must be equal to dimension {axis} of the operand shape {list(aval_in.shape)}.","messagePattern":"Sum of sizes (.+?) must be equal to dimension (.+?) of the operand shape (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":1944,"sourceCode":"\n\n@register_pull_block_spec_rule(lax.split_p)\ndef _split_pull_rule(\n    ctx: PullRuleContext,\n    out_block_transforms: tuple[BlockIndexTransform, ...],\n    *,\n    sizes: Sequence[int],\n    axis: int,\n):\n  aval_in = ctx.avals_in[0]\n  assert isinstance(aval_in, core.ShapedArray)\n  assert all(isinstance(aval, core.ShapedArray) for aval in ctx.avals_out)\n\n  # turn numpy ints into ints\n  sizes = [int(s) for s in sizes]\n  n = sum(sizes)\n  if n != aval_in.shape[axis]:\n    raise NotImplementedError(\n        f'Sum of sizes {n} must be equal to dimension {axis} of the operand '\n        f'shape {list(aval_in.shape)}.'\n    )\n  valid_transforms = [\n      bt for bt in out_block_transforms if bt is not no_block_index_transform\n  ]\n  if not valid_transforms:\n    return [no_block_index_transform]\n\n  block_transform = valid_transforms[0]\n\n  new_block_shape = list(block_transform.block_shape)\n  new_block_shape[axis] = pallas_core.Blocked(n)\n\n  def new_block_index_transform(*idxs):\n    idx = list(block_transform.block_index_transform(*idxs))\n    idx[axis] = 0\n    return tuple(idx)","sourceCodeStart":1926,"sourceCodeEnd":1962,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L1926-L1962","documentation":"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].","triggerScenarios":"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).","commonSituations":"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.","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"],"exampleFix":"// before\nparts = lax.split(x, [2])  # sizes=[2] but x.shape[0] == 5\n// after\nparts = lax.split(x, [2, 5])  # sizes sum (2+3) == x.shape[0]\n","handlingStrategy":"validation","validationCode":"sizes = [int(s) for s in sizes]\nassert sum(sizes) == x.shape[axis], f'{sum(sizes)} != {x.shape[axis]}'","typeGuard":"def split_sizes_valid(x, sizes, axis) -> bool:\n    return sum(int(s) for s in sizes) == x.shape[axis]","tryCatchPattern":"try:\n    parts = fused_split(x)\nexcept NotImplementedError as e:\n    if 'Sum of sizes' in str(e):\n        parts = jnp.split(x, indices_or_sections=len(x.shape[axis] and sizes), axis=axis)\n    else:\n        raise","preventionTips":["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"],"tags":["jax","pallas","lax-split","shape-mismatch","validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}