jax-ml/jax · error · ValueError

reduceat only supported for binary ufuncs

Error message

reduceat only supported for binary ufuncs

What it means

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.

Source

Thrown at jax/_src/numpy/ufunc_api.py:515

      >>> indices = jnp.array([0, 2, 5])
      >>> jnp.add.reduce(x, indices)
      Array([ 3, 12, 21], dtype=int32)

      This is more-or-less equivalent to the following:

      >>> jnp.array([x[0:2].sum(), x[2:5].sum(), x[5:].sum()])
      Array([ 3, 12, 21], dtype=int32)

      For some binary ufuncs, JAX provides similar APIs within :mod:`jax.ops`.
      For example, :meth:`jax.add.reduceat` is similar to :func:`jax.ops.segment_sum`,
      although in this case the segments are defined via an array of segment ids:

      >>> segments = jnp.array([0, 0, 1, 1, 1, 2, 2, 2])
      >>> jax.ops.segment_sum(x, segments)
      Array([ 3, 12, 21], dtype=int32)
    """
    if self.nin != 2:
      raise ValueError("reduceat only supported for binary ufuncs")
    if self.nout != 1:
      raise ValueError("reduceat only supported for functions returning a single value")
    if out is not None:
      raise NotImplementedError(f"out argument of {self.__name__}.reduceat()")

    reduceat = self.__static_props['reduceat'] or self._reduceat_via_scan
    return reduceat(a, indices, axis=axis, dtype=dtype)

  def _reduceat_via_scan(self, a: ArrayLike, indices: Any, axis: int = 0,
                         dtype: DTypeLike | None = None) -> Array:
    check_arraylike(f"{self.__name__}.reduceat", a, indices)
    a = lax.asarray(a)
    idx_tuple = indexing.eliminate_deprecated_list_indexing(indices)
    assert len(idx_tuple) == 1
    indices = idx_tuple[0]
    if a.ndim == 0:
      raise ValueError(f"reduceat: a must have 1 or more dimension, got {a.shape=}")
    if indices.ndim != 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a binary ufunc such as add, multiply, or bitwise_or with .reduceat
2. For custom segmented ops use jax.ops.segment_sum/segment_max or lax.scan directly

Example fix

// before
jnp.negative.reduceat(x, idx)
// after
jnp.add.reduceat(x, idx)

assert ufunc.nin == 2, f'{ufunc.__name__}.reduceat needs a binary ufunc'

def is_binary_ufunc(u): return u.nin == 2

When it happens

Trigger: jnp.negative.reduceat(x, indices) or any unary ufunc's .reduceat.

Common situations: Generic segmented-reduction helper code that accepts an arbitrary ufunc; porting numpy reduceat usage with the wrong op.

Related errors

• reduce only supported for binary ufuncs
• reduce only supported for functions returning a single value
• accumulate only supported for binary ufuncs
• accumulate only supported for functions returning a single v
• {name} wrapped function must be passed at least one argument

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