jax-ml/jax · error · NotImplementedError

where argument of {self}

Error message

where argument of {self}

What it means

JAX ufunc objects accept numpy's where keyword for signature compatibility but vectorized masking (where=) is not implemented for generic ufuncs. Passing a non-None where raises NotImplementedError.

Source

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

    # because they are considered implementation details rather than
    # necessary parts of object identity.
    return hash((self._func, self.__name__, self.identity,
                 self.nin, self.nout, self.nargs))

  def __eq__(self, other: Any) -> bool:
    return isinstance(other, ufunc) and (
      (self._func, self.__name__, self.identity, self.nin, self.nout, self.nargs) ==
      (other._func, other.__name__, other.identity, other.nin, other.nout, other.nargs))

  def __repr__(self) -> str:
    return f"<jnp.ufunc '{self.__name__}'>"

  def __call__(self, *args: ArrayLike, out: None = None, where: None = None) -> Any:
    check_arraylike(self.__name__, *args)
    if out is not None:
      raise NotImplementedError(f"out argument of {self}")
    if where is not None:
      raise NotImplementedError(f"where argument of {self}")
    call = (self.__static_props['call']
            or cast(Callable[..., Any], self._call_vectorized))
    return call(*args)

  @api.jit(static_argnames=['self'])
  def _call_vectorized(self, *args):
    return vectorize(self._func)(*args)

  @api.jit(static_argnames=['self', 'axis', 'dtype', 'out', 'keepdims'])
  def reduce(self, a: ArrayLike, axis: int | None = 0,
             dtype: DTypeLike | None = None,
             out: None = None, keepdims: bool = False, initial: ArrayLike | None = None,
             where: ArrayLike | None = None) -> Array:
    """Reduction operation derived from a binary function.

    JAX implementation of :meth:`numpy.ufunc.reduce`.

    Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace with jnp.where(mask, x op y, default)
2. Use x.at[...].set(...) with boolean masks for conditional assignment

Example fix

// before
jnp.multiply(x, y, where=mask)
// after
jnp.where(mask, x * y, x)

if where is not None:
    result = jnp.where(where, ufunc(*args), default)

When it happens

Trigger: jnp.multiply(x, y, where=mask) on a generic jnp.ufunc.

Common situations: Porting numpy code that uses where= for conditional elementwise updates.

Related errors

• out argument of {self}
• reduction operation {self.__name__!r} does not have an ident
• where argument must have dtype=bool; got dtype={lax._dtype(w
• array() takes at most 5 positional arguments but {len(args)
• array() got multiple values for argument '{name}'

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