jax-ml/jax · error · NotImplementedError

In order to best JIT compile `mode`, we cannot know whether

Error message

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.

What it means

jax.scipy.stats.mode with nan_policy='raise' raises NotImplementedError because detecting NaNs would require a data-dependent branch, which breaks JIT tracing. The message asks users to check for NaNs outside the function.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove nan_policy='raise' and check for NaNs before the call: has_nan = bool(jnp.isnan(x).any()); raise manually if true.
Validate input cleanliness upstream (assert not jnp.isnan(x).any()) and then call mode with 'propagate'.
Use scipy.stats.mode outside JIT when 'raise' semantics are needed.

Example fix

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

// after
if jnp.isnan(x).any():
    raise ValueError("x contains NaNs")
jax.scipy.stats.mode(x, nan_policy='propagate')

Defensive patterns

Strategy: validation

Validate before calling

if jnp.isnan(x).any():
    raise ValueError("x contains NaNs")
jax.scipy.stats.mode(x, nan_policy='propagate')

Prevention

Do NaN checks eagerly outside jitted functions.
Avoid nan_policy='raise' with jax; it always raises.

When it happens

Trigger: Calling jax.scipy.stats.mode(a, nan_policy='raise') — even when the array contains no NaNs, this always raises.

Common situations: Porting scipy defaults (scipy used 'propagate' historically but users often set 'raise'); wrapping mode in jit and wanting NaN safety.

Related errors

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