jax-ml/jax · error · ValueError

where argument must have dtype=bool; got dtype={lax._dtype(w

Error message

where argument must have dtype=bool; got dtype={lax._dtype(where)}

What it means

ufunc.reduce's where mask selects which elements participate in the reduction and must be a boolean array. Passing a non-bool mask (e.g. int 0/1 array or float mask) raises this ValueError reporting the offending dtype.

Source

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

      >>> jnp.bitwise_or.reduce(x, axis=1)
      Array([3, 7], dtype=int32)
    """
    check_arraylike(f"{self.__name__}.reduce", a)
    if self.nin != 2:
      raise ValueError("reduce only supported for binary ufuncs")
    if self.nout != 1:
      raise ValueError("reduce only supported for functions returning a single value")
    if out is not None:
      raise NotImplementedError(f"out argument of {self.__name__}.reduce()")
    if initial is not None:
      check_arraylike(f"{self.__name__}.reduce", initial)
    if where is not None:
      check_arraylike(f"{self.__name__}.reduce", where)
      if self.identity is None and initial is None:
        raise ValueError(f"reduction operation {self.__name__!r} does not have an identity, "
                         "so to use a where mask one has to specify 'initial'.")
      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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the mask: where=mask.astype(bool)
2. Create masks with comparison operators (x > 0) which naturally yield bool

Example fix

// before
jnp.add.reduce(x, where=jnp.array([1,0,1]))
// after
jnp.add.reduce(x, where=jnp.array([1,0,1], dtype=bool))

where = where.astype(bool) if where is not None else where

def is_bool_mask(m): return m is None or lax.dtype(m) == jnp.bool_

When it happens

Trigger: jnp.add.reduce(x, where=jnp.array([1,0,1])) — an int32/int64 mask instead of bool.

Common situations: Masks loaded from files or produced by arithmetic comparisons of integers; numpy code where int masks were implicitly truthy.

Related errors

• reduction operation {self.__name__!r} does not have an ident
• jnp.{name}: where must be None or a boolean array; got {wher
• where argument of {self}
• reduce only supported for binary ufuncs
• reduce only supported for functions returning a single value

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