jax-ml/jax · error · ValueError

jax.numpy.{op} with a symbolic number of sections is not sup

Error message

jax.numpy.{op} with a symbolic number of sections is not supported

What it means

jnp.split-family functions accept either an index list or an integer number of sections, but the integer form must be a concrete (statically known) value. A symbolic dimension (e.g. from jax.export or dynamic shapes) as the number of sections cannot be compiled.

Source

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

    # and the result is clipped to [0, size] so that out-of-bound indices yield
    # empty sections rather than negative sizes. Symbolic indices are left
    # untouched, and clipping is skipped for symbolic sizes, since neither
    # comparison is well-defined for them.
    def _resolve(i_s):
      i = core.concrete_dim_or_error(i_s, f"in jax.numpy.{op} argument 1")
      if core.is_symbolic_dim(i):
        return i
      if i < 0:
        i += size
      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]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass an explicit list of split indices instead of a section count: jnp.split(x, [i * step for i in range(1, n)]) with concrete ints
Make the number of sections a concrete Python int (compute it outside jit/export)
If size is symbolic too, compute split sizes as symbolic expressions via the index-list path

Example fix

// before
jnp.split(x, num_parts)  # num_parts is a symbolic dim
// after
jnp.split(x, [k * (x.shape[0] // num_parts_) for k in range(1, num_parts_)])  # concrete ints

Defensive patterns

Strategy: fallback

Validate before calling

if not isinstance(indices_or_sections, int):
    indices_or_sections = int(indices_or_sections)  # or build index list

Prevention

Pass explicit split index lists under dynamic shapes/jax.export
Compute section counts as concrete Python ints outside jit

When it happens

Trigger: Calling jnp.split(x, n) where n is a symbolic dimension or tracer, e.g. splitting along an axis whose size or divisor comes from export-shaped inputs with dynamic dimensions.

Common situations: Using jax.export / dynamic shapes; passing a lazily-computed section count that is a DimVar rather than a Python int; splitting under transformations that lose concretization.

Related errors

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