jax-ml/jax · error · ValueError

Sizes passed to split must be nonnegative, got {list(sizes)}

Error message

Sizes passed to split must be nonnegative, got {list(sizes)}

What it means

jax.lax.split takes an explicit list of section sizes along an axis; every size must be >= 0. A negative size is meaningless (it would produce a negative dimension) and raises this ValueError immediately in the shape rule.

Source

Thrown at jax/_src/lax/lax.py:7541

unstack_p = core.Primitive('unstack')
unstack_p.multiple_results = True
unstack_p.def_abstract_eval(
    partial(standard_multi_result_abstract_eval, unstack_p, _unstack_shape_rule,
            _unstack_dtype_rule, _unstack_weak_type_rule, _unstack_sharding_rule,
            _unstack_vma_rule, _unstack_ur_rule, None))
unstack_p.def_impl(partial(dispatch.apply_primitive, unstack_p))
ad.deflinear2(unstack_p, _unstack_transpose_rule)
mlir.register_lowering(unstack_p, _unstack_lower)

batching.primitive_batchers[stack_p] = _stack_batch_rule
batching.primitive_batchers[unstack_p] = _unstack_batch_rule


def _split_shape_rule(operand, *, sizes, axis):
  shapes = []
  shape = list(operand.shape)
  if any(s < 0 for s in sizes):
    raise ValueError(
      f"Sizes passed to split must be nonnegative, got {list(sizes)}")
  if operand.shape[axis] != np.sum(sizes):
    raise ValueError(
      f"Sum of sizes {np.sum(sizes)} must be equal to dimension {axis} of the "
      f"operand shape {list(operand.shape)}")
  for size in sizes:
    shape[axis] = size
    shapes.append(tuple(shape))
  return shapes

def _split_dtype_rule(operand, *, sizes, axis):
  return (operand.dtype,) * len(sizes)

def _split_weak_type_rule(operand, *, sizes, axis):
  return (operand.weak_type,) * len(sizes)

def _split_transpose_rule(cotangents, operand, *, sizes, axis):
  assert ad.is_undefined_primal(operand)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Validate sizes before calling: assert all(s >= 0 for s in sizes)
2. Compute the last chunk as the remainder dim - sum(others) and check it's >= 0
3. Use jnp.split / array_split with integer section counts if you don't need explicit sizes

Example fix

# before
parts = jax.lax.split(x, sizes=[k, dim - 2*k])  # negative if 2*k > dim
# after
rest = dim - 2*k
assert rest >= 0, (dim, k)
parts = jax.lax.split(x, sizes=[k, k, rest])

assert all(s >= 0 for s in sizes), sizes

def valid_sizes(sizes) -> bool:
    return all(s >= 0 for s in sizes)

When it happens

Trigger: jax.lax.split(x, sizes=[2, -1, 3]); computing sizes by subtraction that goes negative (e.g. total - used where used > total).

Common situations: Splitting a tensor into fixed + remainder chunks computed as dim - k when k exceeds dim; off-by-one in size arithmetic; sizes derived from config with bad values.

Related errors

• {full_name} must be a pytree prefix with bool leaves or a tu
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin
• `unreduced_kind` should be `None` when `unreduced` is an emp

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