jax-ml/jax · error · NotImplementedError

Logic for `nan_policy` of {nan_policy} is not implemented

Error message

Logic for `nan_policy` of {nan_policy} is not implemented

What it means

jax.scipy.stats.mode accepts nan_policy='omit' in its signature check but has not implemented the omit logic, so it raises NotImplementedError with a TODO. Only nan_policy='propagate' is actually functional in JAX.

Source

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

    >>> 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:
    if axis is None:
      output_shape = tuple(1 for i in input_shape)
    else:
      output_shape = tuple(1 if i == axis else s for i, s in enumerate(input_shape))
  else:
    if axis is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Filter NaNs out of the array yourself before calling mode (e.g. x[~jnp.isnan(x)] for 1-D input).
  2. Call scipy.stats.mode on the host with nan_policy='omit' instead.
  3. File/check the upstream JAX issue and fall back to a manual nan-aware mode computation.

Example fix

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

// after
clean = x[~jnp.isnan(x)]
jax.scipy.stats.mode(clean, nan_policy='propagate')
Defensive patterns

Strategy: fallback

Validate before calling

# pre-clean NaNs instead of nan_policy='omit'
x = x[~jnp.isnan(x)] if x.ndim == 1 else x

Try / catch

try:
    jax.scipy.stats.mode(x, nan_policy='omit')
except NotImplementedError:
    clean = x[~jnp.isnan(x)]
    jax.scipy.stats.mode(clean)

Prevention

When it happens

Trigger: Calling jax.scipy.stats.mode(a, nan_policy='omit').

Common situations: Porting scipy.stats.mode code that used 'omit' to skip NaNs in survey/sensor data; expecting parity with scipy.

Related errors

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