jax-ml/jax · error · ValueError
{nan_policy} is not supported
Error message
{nan_policy} is not supported What it means
jax.scipy.stats.sem (standard error of the mean) only supports nan_policy values 'propagate' and 'omit'. Any other string raises this error. scipy also supports 'raise', which JAX does not implement because JAX errors must be shape/dtype-static under jit.
Source
Thrown at jax/_src/scipy/stats/_core.py:311
... jax.scipy.stats.sem(x2)
Array([1.73, nan, 1.53, nan, nan, nan], dtype=float32)
If ``nan_policy='omit```, ``sem`` omits the ``nan`` values and computes the error
for the remaining values along the specified axis.
>>> with jnp.printoptions(precision=2, suppress=True):
... jax.scipy.stats.sem(x2, nan_policy='omit')
Array([1.73, 1.5 , 1.53, 2. , 2.5 , 0.5 ], dtype=float32)
"""
b, = promote_args_inexact("sem", a)
if nan_policy == "propagate":
size = b.size if axis is None else b.shape[axis]
return b.std(axis, ddof=ddof, keepdims=keepdims) / jnp.sqrt(size).astype(b.dtype)
elif nan_policy == "omit":
count = (~jnp.isnan(b)).sum(axis, keepdims=keepdims)
return jnp.nanstd(b, axis, ddof=ddof, keepdims=keepdims) / jnp.sqrt(count).astype(b.dtype)
else:
raise ValueError(f"{nan_policy} is not supported")
View on GitHub (pinned to 1e1c6a8fc0)
Solutions
- Use nan_policy='propagate' (default) or 'omit'
- If you need fail-on-NaN behavior, add an explicit jnp.isnan(x).any() check with a Python-side raise before calling sem
- Fix typos in the nan_policy string
Example fix
// before se = jax.scipy.stats.sem(x, nan_policy='raise') // after assert not jnp.isnan(x).any() se = jax.scipy.stats.sem(x, nan_policy='propagate')
Defensive patterns
Strategy: validation
Validate before calling
nan_policy = 'omit'
assert nan_policy in ('propagate', 'omit') Type guard
def is_supported_nan_policy(p: str) -> bool:
return p in ('propagate', 'omit') Try / catch
try:
sem = jax.scipy.stats.sem(x, nan_policy=p)
except ValueError:
sem = jax.scipy.stats.sem(x, nan_policy='propagate') Prevention
- Map scipy's 'raise' to an explicit isnan pre-check
- Validate nan_policy from config before passing through
- Remember JAX errors cannot be data-dependent under jit
When it happens
Trigger: Calling jax.scipy.stats.sem(x, nan_policy='raise') or an arbitrary/misspelled string like 'ignore' or 'omit_nan'.
Common situations: Copying scipy.stats.sem code that uses nan_policy='raise'; assuming full scipy parameter parity in jax.scipy.
Related errors
- ind must be a positive integer; got {ind=}
- tensorinv is only possible when the product of the first `in
- Expected kind to be on of: {valid_kind}; got {kind}
- Expected kind to be one of: {valid_kind}; got {kind}
- n must be a positive power of 2; got {n}.
AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27).
Data as JSON: /api/errors/f7bbbf48883f2e29.
Report an issue: GitHub.