{"record":{"id":"2323ec4df1b2ef13","repo":"jax-ml/jax","slug":"array-split-does-not-result-in-an-equal-division","errorCode":null,"errorMessage":"array split does not result in an equal division: rest is {r}","messagePattern":"array split does not result in an equal division: rest is (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":3168,"sourceCode":"      if core.is_symbolic_dim(size):\n        return i\n      return np.clip(i, 0, size)\n    split_indices = np.asarray(\n        [0, *(_resolve(i_s) for i_s in indices_or_sections), size])\n    sizes = list(np.diff(split_indices))\n  else:\n    if core.is_symbolic_dim(indices_or_sections):\n      raise ValueError(f\"jax.numpy.{op} with a symbolic number of sections is \"\n                       \"not supported\")\n    num_sections: int = core.concrete_or_error(int, indices_or_sections,\n                                               f\"in jax.numpy.{op} argument 1\")\n    part_size, r = divmod(size, num_sections)\n    if r == 0:\n      sizes = [part_size] * num_sections\n    elif op == \"array_split\":\n      sizes = [(part_size + 1)] * r + [part_size] * (num_sections - r)\n    else:\n      raise ValueError(f\"array split does not result in an equal division: rest is {r}\")\n  sizes = [i if core.is_symbolic_dim(i) else np.int64(i)\n           for i in sizes]\n  return list(lax.split(ary, sizes, axis=axis))\n\n\n@export\ndef split(ary: ArrayLike, indices_or_sections: int | Sequence[int] | ArrayLike,\n          axis: int = 0) -> list[Array]:\n  \"\"\"Split an array into sub-arrays.\n\n  JAX implementation of :func:`numpy.split`.\n\n  Args:\n    ary: N-dimensional array-like object to split\n    indices_or_sections: either a single integer or a sequence of indices.\n\n      - if ``indices_or_sections`` is an integer *N*, then *N* must evenly divide\n        ``ary.shape[axis]`` and ``ary`` will be divided into *N* equally-sized","sourceCodeStart":3150,"sourceCodeEnd":3186,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L3150-L3186","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use jnp.array_split, which handles uneven splits","Pad or trim the array so size % num_sections == 0 before splitting","Compute sections from size: num = size // chunk if splitting by chunk size"],"exampleFix":"// before\nparts = jnp.split(x, 3)  # x.shape[0] == 10\n// after\nparts = jnp.array_split(x, 3)","handlingStrategy":"fallback","validationCode":"import jax.numpy as jnp\nif size % num_sections != 0:\n    parts = jnp.array_split(x, num_sections)\nelse:\n    parts = jnp.split(x, num_sections)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Default to array_split when divisibility isn't guaranteed","Pad batches to a multiple of world_size before splitting"],"tags":["jax","split","uneven-division"],"backgroundTag":"uneven-split-division","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}