jax-ml/jax · error · ValueError

Illegal nan_policy value {nan_policy!r}; expected one of {'p

Error message

Illegal nan_policy value {nan_policy!r}; expected one of {'propagate', 'omit', 'raise'}

What it means

jax.scipy.stats.mode validates its nan_policy keyword and only accepts 'propagate', 'omit', or 'raise'; any other string raises this ValueError. This mirrors scipy's API surface but is enforced eagerly at call time in JAX.

Source

Thrown at jax/_src/scipy/stats/_core.py:91

    (Array([1, 3, 2], dtype=int32), Array([3, 3, 3], dtype=int32))

    By default, ``jax.scipy.stats.mode`` reduces the dimension of the result.
    To keep the dimensions same as that of the input array, the argument
    ``keepdims`` must be set to ``True``.

    >>> mode, count = jax.scipy.stats.mode(x1, axis=1, keepdims=True)
    >>> mode, count
    (Array([[1],
           [3],
           [2]], dtype=int32), Array([[3],
           [3],
           [3]], dtype=int32))
  """
  check_arraylike("mode", a)
  x = jnp.atleast_1d(a)

  if nan_policy not in ["propagate", "omit", "raise"]:
    raise ValueError(
      f"Illegal nan_policy value {nan_policy!r}; expected one of "
      "{'propagate', 'omit', 'raise'}"
    )
  if nan_policy == "omit":
    # TODO: return answer without nans included.
    raise NotImplementedError(
      f"Logic for `nan_policy` of {nan_policy} is not implemented"
    )
  if nan_policy == "raise":
    raise NotImplementedError(
      "In order to best JIT compile `mode`, we cannot know whether `x` contains nans. "
      "Please check if nans exist in `x` outside of the `mode` function."
    )
  if axis is not None:
    axis = canonicalize_axis(axis, x.ndim)

  input_shape = x.shape
  if keepdims:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use one of the exact strings 'propagate', 'omit', or 'raise' (lowercase).
Note that 'omit' and 'raise' themselves raise NotImplementedError in jax — effectively only 'propagate' works; prefer removing the kwarg or handling NaNs before the call.

Example fix

// before
jax.scipy.stats.mode(x, nan_policy='skip')

// after
jax.scipy.stats.mode(x, nan_policy='propagate')

Defensive patterns

Strategy: validation

Validate before calling

assert nan_policy in ('propagate', 'omit', 'raise'), f"bad nan_policy: {nan_policy!r}"

Type guard

def is_valid_nan_policy(p) -> bool:
    return isinstance(p, str) and p in ('propagate', 'omit', 'raise')

Try / catch

try:
    jax.scipy.stats.mode(x, nan_policy=p)
except ValueError as e:
    if 'Illegal nan_policy' in str(e):
        jax.scipy.stats.mode(x)  # fallback to default
    else:
        raise

Prevention

Centralize nan_policy constants instead of free-form strings.
Remember only 'propagate' is implemented in jax.

When it happens

Trigger: Calling jax.scipy.stats.mode(a, nan_policy='skip') or a typo like 'raises', 'Propagate', or None.

Common situations: Copy-pasted scipy code using older/newer policy names; passing an uninitialized variable (None) as nan_policy.

Related errors

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