{"record":{"id":"6da9a2c1d4df75bd","repo":"jax-ml/jax","slug":"sum-of-sizes-np-sum-sizes-must-be-equal-to-dime","errorCode":null,"errorMessage":"Sum of sizes {np.sum(sizes)} must be equal to dimension {axis} of the operand shape {list(operand.shape)}","messagePattern":"Sum of sizes (.+?) must be equal to dimension (.+?) of the operand shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":7544,"sourceCode":"    partial(standard_multi_result_abstract_eval, unstack_p, _unstack_shape_rule,\n            _unstack_dtype_rule, _unstack_weak_type_rule, _unstack_sharding_rule,\n            _unstack_vma_rule, _unstack_ur_rule, None))\nunstack_p.def_impl(partial(dispatch.apply_primitive, unstack_p))\nad.deflinear2(unstack_p, _unstack_transpose_rule)\nmlir.register_lowering(unstack_p, _unstack_lower)\n\nbatching.primitive_batchers[stack_p] = _stack_batch_rule\nbatching.primitive_batchers[unstack_p] = _unstack_batch_rule\n\n\ndef _split_shape_rule(operand, *, sizes, axis):\n  shapes = []\n  shape = list(operand.shape)\n  if any(s < 0 for s in sizes):\n    raise ValueError(\n      f\"Sizes passed to split must be nonnegative, got {list(sizes)}\")\n  if operand.shape[axis] != np.sum(sizes):\n    raise ValueError(\n      f\"Sum of sizes {np.sum(sizes)} must be equal to dimension {axis} of the \"\n      f\"operand shape {list(operand.shape)}\")\n  for size in sizes:\n    shape[axis] = size\n    shapes.append(tuple(shape))\n  return shapes\n\ndef _split_dtype_rule(operand, *, sizes, axis):\n  return (operand.dtype,) * len(sizes)\n\ndef _split_weak_type_rule(operand, *, sizes, axis):\n  return (operand.weak_type,) * len(sizes)\n\ndef _split_transpose_rule(cotangents, operand, *, sizes, axis):\n  assert ad.is_undefined_primal(operand)\n  if all(type(t) is ad_util.Zero for t in cotangents):\n    return [ad_util.Zero(operand.aval)]\n  cotangents = [ct.instantiate() if type(ct) is ad_util.Zero else ct","sourceCodeStart":7526,"sourceCodeEnd":7562,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L7526-L7562","documentation":"jax.lax.split requires the sizes to exactly tile the chosen axis: sum(sizes) must equal operand.shape[axis]. This ValueError reports the computed sum and the actual dimension so you can see the mismatch.","triggerScenarios":"jax.lax.split(x, sizes=[2,3]) on an axis of length 6; hardcoded sizes after the input shape changed; remainder chunk computed with off-by-one.","commonSituations":"Dataset/tensor width changed (vocab size, feature dim) while split sizes stayed hardcoded; splitting sequence length into chunks that don't divide evenly.","solutions":["Derive the last size from the axis: sizes[-1] = x.shape[axis] - sum(sizes[:-1]) and assert >= 0","Use jnp.array_split or jnp.split for equal chunks when exact tiling is acceptable","Parameterize sizes from the runtime shape instead of constants"],"exampleFix":"# before\nparts = jax.lax.split(x, sizes=[128, 128, 8])  # axis len changed to 264\n# after\nrest = x.shape[-1] - 256\nparts = jax.lax.split(x, sizes=[128, 128, rest])","handlingStrategy":"validation","validationCode":"assert x.shape[axis] == sum(sizes), (x.shape[axis], sizes)\n# or derive: sizes = sizes[:-1] + [x.shape[axis] - sum(sizes[:-1])]","typeGuard":"def sizes_tile_axis(sizes, axis_len) -> bool:\n    return sum(sizes) == axis_len","tryCatchPattern":null,"preventionTips":["Derive the last split size from the runtime axis length","Prefer jnp.array_split for uneven chunks"],"tags":["jax","split","size-mismatch","shape-validation"],"backgroundTag":"invalid-split-sizes","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}