jax-ml/jax · error · NotImplementedError

tuple of axes

Error message

tuple of axes

What it means

In the scan-based fallback implementation of ufunc.reduce, reducing over a tuple of multiple axes at once is not implemented. Note the raise is unreachable in normal flow (it follows the tuple construction), but tuple-axis reduction on generic ufuncs without a registered reduce is unsupported.

Source

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

      if lax._dtype(where) != bool:
        raise ValueError(f"where argument must have dtype=bool; got dtype={lax._dtype(where)}")
    reduce = self.__static_props['reduce'] or self._reduce_via_scan
    return reduce(a, axis=axis, dtype=dtype, keepdims=keepdims, initial=initial, where=where)

  def _reduce_via_scan(self, arr: ArrayLike, axis: int | tuple[int, ...] | None = 0, dtype: DTypeLike | None = None,
                       keepdims: bool = False, initial: ArrayLike | None = None,
                       where: ArrayLike | None = None) -> Array:
    assert self.nin == 2 and self.nout == 1
    arr = lax.asarray(arr)
    if initial is None:
      initial = self.identity
    if dtype is None:
      dtype = api.eval_shape(self._func, lax._one(arr), lax._one(arr)).dtype
    if where is not None:
      where = _broadcast_to(where, arr.shape)
    if isinstance(axis, tuple):
      axis = tuple(canonicalize_axis(a, arr.ndim) for a in axis)
      raise NotImplementedError("tuple of axes")
    elif axis is None:
      if keepdims:
        final_shape = (1,) * arr.ndim
      else:
        final_shape = ()
      arr = arr.ravel()
      if where is not None:
        where = where.ravel()
      axis = 0
    else:
      axis = canonicalize_axis(axis, arr.ndim)
      if keepdims:
        final_shape = (*arr.shape[:axis], 1, *arr.shape[axis + 1:])
      else:
        final_shape = (*arr.shape[:axis], *arr.shape[axis + 1:])

    # TODO: handle without transpose?
    if axis != 0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce one axis at a time in a loop
2. Prefer jnp.sum/jnp.prod/jnp.bitwise_or.reduce which have native multi-axis support
3. Flatten the array first with .reshape(-1) when reducing over all axes

Example fix

// before
jnp.logical_or.reduce(x, axis=(0,1))
// after
jnp.logical_or.reduce(x.ravel())

if isinstance(axis, tuple):
    for a in axis:
        x = ufunc.reduce(x, axis=a)

When it happens

Trigger: Calling .reduce with axis=(0,1) on a ufunc that lacks a static reduce implementation, falling through to _reduce_via_scan.

Common situations: Using exotic ufuncs (only the scan fallback exists) together with multi-axis reduction.

Related errors

• reduce only supported for binary ufuncs
• reduce only supported for functions returning a single value
• out argument of {self.__name__}.reduce()
• reduction operation {self.__name__!r} does not have an ident
• where argument must have dtype=bool; got dtype={lax._dtype(w

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