jax-ml/jax · error · ValueError

num_segments must be non-negative.

Error message

num_segments must be non-negative.

What it means

Raised by JAX's segment reduction ops (segment_sum/prod/max/min) when the num_segments argument, after being resolved to a concrete value, is negative. num_segments determines the size of the output array, so a negative value is meaningless. JAX validates it eagerly in _segment_update before allocating the output.

Source

Thrown at jax/_src/ops/scatter.py:195

                    segment_ids: ArrayLike,
                    scatter_op: Callable,
                    num_segments: int | None = None,
                    indices_are_sorted: bool = False,
                    unique_indices: bool = False,
                    bucket_size: int | None = None,
                    reducer: Callable | None = None,
                    mode: slicing.GatherScatterMode | str | None = None,
                    out_sharding: NamedSharding | None = None) -> Array:
  check_arraylike(name, data, segment_ids)
  mode = slicing.GatherScatterMode.FILL_OR_DROP if mode is None else mode
  data = jnp.asarray(data)
  segment_ids = jnp.asarray(segment_ids)
  dtype = data.dtype
  if num_segments is None:
    num_segments = np.max(segment_ids) + 1
  num_segments = core.concrete_dim_or_error(num_segments, "segment_sum() `num_segments` argument.")
  if num_segments is not None and num_segments < 0:
    raise ValueError("num_segments must be non-negative.")

  if bucket_size is None:
    out = jnp.full((num_segments,) + data.shape[1:],
                   _get_identity(scatter_op, dtype), dtype=dtype)
    return _scatter_update(
      out, segment_ids, data, scatter_op, indices_are_sorted,
      unique_indices, normalize_indices=False, mode=mode,
      out_sharding=out_sharding)

  # Bucketize indices and perform segment_update on each bucket to improve
  # numerical stability for operations like product and sum.
  assert reducer is not None
  if out_sharding is not None:
    raise NotImplementedError
  num_buckets = util.ceil_of_ratio(segment_ids.size, bucket_size)
  out = jnp.full((num_buckets, num_segments) + data.shape[1:],
                 _get_identity(scatter_op, dtype), dtype=dtype)
  out = _scatter_update(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check how num_segments is computed — most often it's max(segment_ids)+1 or a count expression with a sign/off-by-one bug
2. Pass num_segments=None to let JAX infer it as np.max(segment_ids)+1
3. Guard with max(0, num_segments) if negative values can legitimately occur in your pipeline
4. Print/inspect the concrete value before the call when under jit (static_argnums)

Example fix

// before
out = jax.ops.segment_sum(data, ids, num_segments=n_segments - 1)  # n_segments==0
// after
out = jax.ops.segment_sum(data, ids, num_segments=max(0, n_segments - 1))

n = int(num_segments)
assert n >= 0, f'num_segments must be >= 0, got {n}'
out = jax.ops.segment_sum(data, ids, num_segments=n)

When it happens

Trigger: Calling jax.lax.segment_sum/segment_prod/segment_max/segment_min (or jax.ops.segment_*) with a negative num_segments, e.g. segment_sum(data, segment_ids, num_segments=-1), or with a traced value that concrete evaluation resolves to a negative number.

Common situations: Computing num_segments as max(segment_ids)+1 minus an offset (e.g. n_classes - offset) where the offset exceeds the count; off-by-one bugs; passing a Python expression that evaluates negative under jit with static_argnums.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• `unroll` must be a `bool` or a non-negative `int`.
• `compute_on`'s compute_type argument must be a string.
• Invalid value "{default}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty
• new enum value must be None or in {enum_values}, got {new_va

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