jax-ml/jax · error · ValueError

jnp.{name}: where must be None or a boolean array; got {wher

Error message

jnp.{name}: where must be None or a boolean array; got {where.dtype=}.

What it means

JAX reductions (sum, mean, var, logsumexp, etc.) require the where mask to be a boolean array (or None). Unlike numpy, JAX will not implicitly cast integer/float masks to bool, because such casts are a common source of silent bugs under jit.

Source

Thrown at jax/_src/numpy/reductions.py:93

  # default dtype as defined by dtypes.int_ or dtypes.uint.
  if dtypes.issubdtype(dtype, np.bool_):
    return dtypes.default_int_dtype()
  elif dtypes.issubdtype(dtype, np.unsignedinteger):
    default_uint_dtype = dtypes.default_uint_dtype()
    if np.iinfo(dtype).bits < np.iinfo(default_uint_dtype).bits:
      return default_uint_dtype
  elif dtypes.issubdtype(dtype, np.integer):
    default_int_dtype = dtypes.default_int_dtype()
    if np.iinfo(dtype).bits < np.iinfo(default_int_dtype).bits:
      return default_int_dtype
  return dtype

def check_where(name: str, where: ArrayLike | None) -> Array | None:
  if where is None:
    return where
  where = ensure_arraylike(name, where)
  if where.dtype != bool:
    raise ValueError(
      f"jnp.{name}: where must be None or a boolean array; got {where.dtype=}."
    )
  return where

ReductionOp = Callable[[Any, Any], Any]

def _reduction(a: ArrayLike, name: str, op: ReductionOp, init_val: ArrayLike,
               *, has_identity: bool = True,
               preproc: Callable[[Array], Array] | None = None,
               bool_op: ReductionOp | None = None,
               upcast_f16_for_computation: bool = False,
               axis: Axis = None, dtype: DTypeLike | None = None, out: None = None,
               keepdims: bool = False, initial: ArrayLike | None = None,
               where_: ArrayLike | None = None,
               parallel_reduce: Callable[..., Array] | None = None,
               promote_integers: bool = False) -> Array:
  bool_op = bool_op or op
  # Note: we must accept out=None as an argument, because numpy reductions delegate to

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the mask: where=mask.astype(bool)
2. Build masks with comparisons: where=x > threshold
3. For weights, multiply the data instead: jnp.sum(a * w) / jnp.sum(w)

Example fix

// before
jnp.mean(a, where=weights)  # weights is float array
// after
jnp.sum(a * weights) / jnp.sum(weights)
# or for a 0/1 mask:
jnp.mean(a, where=mask.astype(bool))

import jax.numpy as jnp
mask = jnp.asarray(mask)
if mask.dtype != jnp.bool_:
    mask = mask.astype(bool)
jnp.mean(a, where=mask)

def is_bool_mask(w) -> bool:
    import jax.numpy as jnp
    return hasattr(w, 'dtype') and jnp.asarray(w).dtype == jnp.bool_

When it happens

Trigger: Passing where=x > 0 works (bool), but where=int_mask or where=float weights array to jnp.sum/mean/var/logsumexp raises this. E.g. jnp.mean(a, where=jnp.array([1,0,1])).

Common situations: Porting numpy code that uses 0/1 integer masks; using a weights array where a mask was expected (should multiply instead); masks produced by arithmetic rather than comparisons.

Related errors

• reduction operation {name} does not have an identity, so to
• integer argument required; got dtype={arr.dtype}
• where argument must have dtype=bool; got dtype={lax._dtype(w
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass

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