jax-ml/jax · error · ValueError

Illegal nan_policy value

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.

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.

Appendix: 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)