jax-ml/jax · error · ValueError

array split does not result in an equal division: rest is {r

Error message

array split does not result in an equal division: rest is {r}

What it means

jnp.split/vsplit/hsplit/dsplit require the axis size to divide evenly by the number of sections (unlike array_split, which tolerates a remainder). If divmod leaves a remainder r, this error is raised.

Source

Thrown at jax/_src/numpy/lax_numpy.py:3168

      if core.is_symbolic_dim(size):
        return i
      return np.clip(i, 0, size)
    split_indices = np.asarray(
        [0, *(_resolve(i_s) for i_s in indices_or_sections), size])
    sizes = list(np.diff(split_indices))
  else:
    if core.is_symbolic_dim(indices_or_sections):
      raise ValueError(f"jax.numpy.{op} with a symbolic number of sections is "
                       "not supported")
    num_sections: int = core.concrete_or_error(int, indices_or_sections,
                                               f"in jax.numpy.{op} argument 1")
    part_size, r = divmod(size, num_sections)
    if r == 0:
      sizes = [part_size] * num_sections
    elif op == "array_split":
      sizes = [(part_size + 1)] * r + [part_size] * (num_sections - r)
    else:
      raise ValueError(f"array split does not result in an equal division: rest is {r}")
  sizes = [i if core.is_symbolic_dim(i) else np.int64(i)
           for i in sizes]
  return list(lax.split(ary, sizes, axis=axis))


@export
def split(ary: ArrayLike, indices_or_sections: int | Sequence[int] | ArrayLike,
          axis: int = 0) -> list[Array]:
  """Split an array into sub-arrays.

  JAX implementation of :func:`numpy.split`.

  Args:
    ary: N-dimensional array-like object to split
    indices_or_sections: either a single integer or a sequence of indices.

      - if ``indices_or_sections`` is an integer *N*, then *N* must evenly divide
        ``ary.shape[axis]`` and ``ary`` will be divided into *N* equally-sized

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jnp.array_split, which handles uneven splits
2. Pad or trim the array so size % num_sections == 0 before splitting
3. Compute sections from size: num = size // chunk if splitting by chunk size

Example fix

// before
parts = jnp.split(x, 3)  # x.shape[0] == 10
// after
parts = jnp.array_split(x, 3)

import jax.numpy as jnp
if size % num_sections != 0:
    parts = jnp.array_split(x, num_sections)
else:
    parts = jnp.split(x, num_sections)

When it happens

Trigger: jnp.split(jnp.arange(10), 3) — 10 % 3 != 0. Common with batch sizes not divisible by the requested split count, or reshaping assumptions that silently broke.

Common situations: Data-parallel sharding where world_size doesn't divide batch size; downstream shape changes (padding removed) breaking previously-even splits; off-by-one in computed section counts.

Related errors

• Sizes passed to split must be nonnegative, got {list(sizes)}
• Sum of sizes {np.sum(sizes)} must be equal to dimension {axi
• jax.numpy.{op} with a symbolic number of sections is not sup
• Only power-of-2 num parts supported.
• Only equal-sized splits are supported.

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