jax-ml/jax · error · ValueError

reduction operation {name} does not have an identity, so to

Error message

reduction operation {name} does not have an identity, so to use a where mask one has to specify 'initial'

What it means

Reductions like max/min have no identity element (no neutral value for the operation), so when a where mask excludes every element along a reduction axis there is no defined result unless you supply an initial value.

Source

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

               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
  # object methods. For example `np.sum(x)` will call `x.sum()` if the `sum()` method
  # exists, passing along all its arguments.
  if out is not None:
    raise NotImplementedError(f"The 'out' argument to jnp.{name} is not supported.")
  a = ensure_arraylike(name, a)
  where_ = check_where(name, where_)
  axis = core.concrete_or_error(None, axis, f"axis argument to jnp.{name}().")

  if initial is None and not has_identity and where_ is not None:
    raise ValueError(f"reduction operation {name} does not have an identity, so to use a "
                     f"where mask one has to specify 'initial'")

  a = preproc(a) if preproc else a
  pos_dims, dims = _reduction_dims(a, axis)

  if initial is None and not has_identity:
    shape = np.shape(a)
    if not _all(shape[d] >= 1 for d in pos_dims):
      raise ValueError(f"zero-size array to reduction operation {name} which has no identity")

  result_dtype: DType
  if dtype is None:
    result_dtype = a.dtype
    if promote_integers:
      result_dtype = _promote_integer_dtype(result_dtype)
  else:
    result_dtype = dtypes.check_and_canonicalize_user_dtype(dtype, name)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Supply an initial value: jnp.max(x, where=mask, initial=-jnp.inf)
  2. Use jnp.where(mask, x, -jnp.inf).max(axis=...) to make the identity explicit
  3. Switch to a reduction with identity (e.g. sum of where-filled values) if semantics allow

Example fix

// before
jnp.max(x, where=mask, axis=1)
// after
jnp.max(x, where=mask, axis=1, initial=-jnp.inf)
Defensive patterns

Strategy: validation

Validate before calling

jnp.max(x, where=mask, initial=-jnp.inf)  # always supply initial with where

Prevention

When it happens

Trigger: Calling jnp.max(x, where=mask) (or min) without initial; if the mask is all False along some axis the result is undefined, so JAX requires initial up front regardless of mask contents.

Common situations: Masked max/min aggregations (e.g. max over valid timesteps); porting numpy which returns garbage/NaN in this case; also triggered under jit where mask contents are unknown so the check is static.

Related errors


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