{"record":{"id":"548a8f1d2381bb49","repo":"jax-ml/jax","slug":"reduceat-only-supported-for-binary-ufuncs","errorCode":null,"errorMessage":"reduceat only supported for binary ufuncs","messagePattern":"reduceat only supported for binary ufuncs","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/ufunc_api.py","lineNumber":515,"sourceCode":"      >>> indices = jnp.array([0, 2, 5])\n      >>> jnp.add.reduce(x, indices)\n      Array([ 3, 12, 21], dtype=int32)\n\n      This is more-or-less equivalent to the following:\n\n      >>> jnp.array([x[0:2].sum(), x[2:5].sum(), x[5:].sum()])\n      Array([ 3, 12, 21], dtype=int32)\n\n      For some binary ufuncs, JAX provides similar APIs within :mod:`jax.ops`.\n      For example, :meth:`jax.add.reduceat` is similar to :func:`jax.ops.segment_sum`,\n      although in this case the segments are defined via an array of segment ids:\n\n      >>> segments = jnp.array([0, 0, 1, 1, 1, 2, 2, 2])\n      >>> jax.ops.segment_sum(x, segments)\n      Array([ 3, 12, 21], dtype=int32)\n    \"\"\"\n    if self.nin != 2:\n      raise ValueError(\"reduceat only supported for binary ufuncs\")\n    if self.nout != 1:\n      raise ValueError(\"reduceat only supported for functions returning a single value\")\n    if out is not None:\n      raise NotImplementedError(f\"out argument of {self.__name__}.reduceat()\")\n\n    reduceat = self.__static_props['reduceat'] or self._reduceat_via_scan\n    return reduceat(a, indices, axis=axis, dtype=dtype)\n\n  def _reduceat_via_scan(self, a: ArrayLike, indices: Any, axis: int = 0,\n                         dtype: DTypeLike | None = None) -> Array:\n    check_arraylike(f\"{self.__name__}.reduceat\", a, indices)\n    a = lax.asarray(a)\n    idx_tuple = indexing.eliminate_deprecated_list_indexing(indices)\n    assert len(idx_tuple) == 1\n    indices = idx_tuple[0]\n    if a.ndim == 0:\n      raise ValueError(f\"reduceat: a must have 1 or more dimension, got {a.shape=}\")\n    if indices.ndim != 1:","sourceCodeStart":497,"sourceCodeEnd":533,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/ufunc_api.py#L497-L533","documentation":"ufunc.reduceat (segmented reduction) requires a binary ufunc (nin == 2). Calling .reduceat on a unary or non-binary ufunc raises this ValueError before any work is done.","triggerScenarios":"jnp.negative.reduceat(x, indices) or any unary ufunc's .reduceat.","commonSituations":"Generic segmented-reduction helper code that accepts an arbitrary ufunc; porting numpy reduceat usage with the wrong op.","solutions":["Use a binary ufunc such as add, multiply, or bitwise_or with .reduceat","For custom segmented ops use jax.ops.segment_sum/segment_max or lax.scan directly"],"exampleFix":"// before\njnp.negative.reduceat(x, idx)\n// after\njnp.add.reduceat(x, idx)","handlingStrategy":"validation","validationCode":"assert ufunc.nin == 2, f'{ufunc.__name__}.reduceat needs a binary ufunc'","typeGuard":"def is_binary_ufunc(u): return u.nin == 2","tryCatchPattern":null,"preventionTips":["Prefer jax.ops.segment_* functions for segmented reductions"],"tags":["jax","ufunc","reduceat","api-misuse"],"backgroundTag":"unsupported-reduce-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}