jax-ml/jax · error · ValueError

Sum of sizes {np.sum(sizes)} must be equal to dimension {axi

Error message

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

What it means

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.

Source

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

    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)
  if all(type(t) is ad_util.Zero for t in cotangents):
    return [ad_util.Zero(operand.aval)]
  cotangents = [ct.instantiate() if type(ct) is ad_util.Zero else ct

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Derive the last size from the axis: sizes[-1] = x.shape[axis] - sum(sizes[:-1]) and assert >= 0
2. Use jnp.array_split or jnp.split for equal chunks when exact tiling is acceptable
3. Parameterize sizes from the runtime shape instead of constants

Example fix

# before
parts = jax.lax.split(x, sizes=[128, 128, 8])  # axis len changed to 264
# after
rest = x.shape[-1] - 256
parts = jax.lax.split(x, sizes=[128, 128, rest])

assert x.shape[axis] == sum(sizes), (x.shape[axis], sizes)
# or derive: sizes = sizes[:-1] + [x.shape[axis] - sum(sizes[:-1])]

def sizes_tile_axis(sizes, axis_len) -> bool:
    return sum(sizes) == axis_len

When it happens

Trigger: 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.

Common situations: Dataset/tensor width changed (vocab size, feature dim) while split sizes stayed hardcoded; splitting sequence length into chunks that don't divide evenly.

Related errors

• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the
• Wrong number of explicit pads for convolution: expected {},

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