{"record":{"id":"cffe6756ee3dbe59","repo":"jax-ml/jax","slug":"num-segments-must-be-non-negative","errorCode":null,"errorMessage":"num_segments must be non-negative.","messagePattern":"num_segments must be non-negative\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/ops/scatter.py","lineNumber":195,"sourceCode":"                    segment_ids: ArrayLike,\n                    scatter_op: Callable,\n                    num_segments: int | None = None,\n                    indices_are_sorted: bool = False,\n                    unique_indices: bool = False,\n                    bucket_size: int | None = None,\n                    reducer: Callable | None = None,\n                    mode: slicing.GatherScatterMode | str | None = None,\n                    out_sharding: NamedSharding | None = None) -> Array:\n  check_arraylike(name, data, segment_ids)\n  mode = slicing.GatherScatterMode.FILL_OR_DROP if mode is None else mode\n  data = jnp.asarray(data)\n  segment_ids = jnp.asarray(segment_ids)\n  dtype = data.dtype\n  if num_segments is None:\n    num_segments = np.max(segment_ids) + 1\n  num_segments = core.concrete_dim_or_error(num_segments, \"segment_sum() `num_segments` argument.\")\n  if num_segments is not None and num_segments < 0:\n    raise ValueError(\"num_segments must be non-negative.\")\n\n  if bucket_size is None:\n    out = jnp.full((num_segments,) + data.shape[1:],\n                   _get_identity(scatter_op, dtype), dtype=dtype)\n    return _scatter_update(\n      out, segment_ids, data, scatter_op, indices_are_sorted,\n      unique_indices, normalize_indices=False, mode=mode,\n      out_sharding=out_sharding)\n\n  # Bucketize indices and perform segment_update on each bucket to improve\n  # numerical stability for operations like product and sum.\n  assert reducer is not None\n  if out_sharding is not None:\n    raise NotImplementedError\n  num_buckets = util.ceil_of_ratio(segment_ids.size, bucket_size)\n  out = jnp.full((num_buckets, num_segments) + data.shape[1:],\n                 _get_identity(scatter_op, dtype), dtype=dtype)\n  out = _scatter_update(","sourceCodeStart":177,"sourceCodeEnd":213,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/ops/scatter.py#L177-L213","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","Pass num_segments=None to let JAX infer it as np.max(segment_ids)+1","Guard with max(0, num_segments) if negative values can legitimately occur in your pipeline","Print/inspect the concrete value before the call when under jit (static_argnums)"],"exampleFix":"// before\nout = jax.ops.segment_sum(data, ids, num_segments=n_segments - 1)  # n_segments==0\n// after\nout = jax.ops.segment_sum(data, ids, num_segments=max(0, n_segments - 1))","handlingStrategy":"validation","validationCode":"n = int(num_segments)\nassert n >= 0, f'num_segments must be >= 0, got {n}'\nout = jax.ops.segment_sum(data, ids, num_segments=n)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Compute num_segments as max(0, int(np.max(segment_ids)) + 1)","Mark num_segments static under jit (static_argnums) so failures are eager and readable"],"tags":["jax","segment-reduction","argument-validation","valueerror"],"backgroundTag":"invalid-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}