jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

JAX uses immutable arrays, so the NumPy out= buffer-write pattern cannot be supported. jnp.std exposes out only for API compatibility and raises NotImplementedError whenever it is not None.

Source

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

  if correction is None:
    correction = ddof
  elif not isinstance(ddof, int) or ddof != 0:
    raise ValueError("ddof and correction can't be provided simultaneously.")
  a = ensure_arraylike("std", a)
  return _std(a, axis=_ensure_optional_axes(axis), dtype=dtype, out=out, correction=correction, keepdims=keepdims,
              where=where, mean=mean)

@api.jit(static_argnames=('axis', 'dtype', 'keepdims'))
def _std(a: Array, *, axis: Axis = None, dtype: DTypeLike | None = None,
         out: None = None, correction: int | float = 0, keepdims: bool = False,
         where: ArrayLike | None = None, mean: ArrayLike | None = None) -> Array:
  where = check_where("std", where)
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "std")
    if not dtypes.issubdtype(dtype, np.inexact):
      raise ValueError(f"dtype argument to jnp.std must be inexact; got {dtype}")
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.std is not supported.")
  return lax.sqrt(var(a, axis=axis, dtype=dtype, correction=correction,
                      keepdims=keepdims, where=where, mean=mean))


@export
def ptp(a: ArrayLike, axis: Axis = None, out: None = None,
        keepdims: bool = False) -> Array:
  r"""Return the peak-to-peak range along a given axis.

  JAX implementation of :func:`numpy.ptp`.

  Args:
    a: input array.
    axis: optional, int or sequence of ints, default=None. Axis along which the
      range is computed. If None, the range is computed on the flattened array.
    keepdims: bool, default=False. If true, reduced axes are left in the result
      with size 1.
    out: Unused by JAX.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Drop the out argument and use the returned array: y = jnp.std(x)
  2. If reusing a buffer name is needed, reassign: buf = jnp.std(x) at runtime (jit will still avoid reallocation)

Example fix

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

Strategy: validation

Validate before calling

kwargs.pop('out', None)  # strip NumPy-style out before calling jnp.std

Prevention

When it happens

Trigger: Calling jnp.std(x, out=buf) or any positional third argument being interpreted as out.

Common situations: Porting NumPy code that uses out= to avoid allocations; generic wrappers forwarding all NumPy kwargs.

Related errors


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