jax-ml/jax · error · NotImplementedError

The 'out' argument to jnp.nanstd is not supported.

Error message

The 'out' argument to jnp.nanstd is not supported.

What it means

jnp.nanstd does not support the NumPy-style out= in-place output argument; it raises NotImplementedError before delegating to nanvar/sqrt. This mirrors JAX's blanket policy of immutable outputs.

Source

Thrown at jax/_src/numpy/reductions.py:2028

    >>> with jnp.printoptions(precision=2, suppress=True):
    ...   print(jnp.nanstd(x, axis=0, keepdims=True, ddof=1))
    [[0.71 0.71 1.41 1.41]]

    To include specific elements of the array to compute standard deviation, you
    can use ``where``.

    >>> where=jnp.array([[1, 0, 1, 0],
    ...                  [0, 1, 0, 1],
    ...                  [1, 1, 0, 1]], dtype=bool)
    >>> jnp.nanstd(x, axis=0, keepdims=True, where=where)
    Array([[0.5, 0.5, 0. , 0. ]], dtype=float32)
  """
  a = ensure_arraylike("nanstd", a)
  where = check_where("nanstd", where)
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "nanstd")
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.nanstd is not supported.")
  return lax.sqrt(nanvar(a, axis=axis, dtype=dtype, ddof=ddof,
                         keepdims=keepdims, where=where, mean=mean))


def _cumulative_reduction(
    name: str, reduction: Callable[..., Array],
    a: ArrayLike, axis: int | None, dtype: DTypeLike | None, out: None = None,
    fill_nan: bool = False, fill_value: ArrayLike = 0,
    promote_integers: bool = False) -> Array:
  """Helper function for implementing cumulative reductions."""
  a = ensure_arraylike(name, a)
  if out is not None:
    raise NotImplementedError(f"The 'out' argument to jnp.{name} is not supported")

  if axis is None or _isscalar(a):
    if not builtins.all(s is None for s in core.typeof(a).sharding.spec):
      raise core.ShardingTypeError(
          "The input should be fully replicated when axis is not specified to"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Remove the out argument and assign the result
  2. Replace out=-based allocation tricks with jax.jit-compiled kernels

Example fix

// before
jnp.nanstd(x, out=buf)
// after
buf = jnp.nanstd(x)
Defensive patterns

Strategy: validation

Validate before calling

s = jnp.nanstd(a, ddof=ddof)  # simply never pass out

Prevention

When it happens

Trigger: Calling jnp.nanstd(a, out=buf) directly.

Common situations: NumPy→JAX ports of NaN-aware standard-deviation computations that relied on out= for memory reuse.

Related errors

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