jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.nanvar rejects the out= argument (NotImplementedError) since JAX cannot write results into a caller-provided mutable buffer. The parameter is retained purely for NumPy API compatibility.

Source

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

     [ 4.  ]]

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

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

def _nanvar(a: Array, *, axis: Axis = None, dtype: DTypeLike | None = None, out: None = None,
           ddof: int = 0, keepdims: bool = False,
           where: ArrayLike | None = None, a_mean: ArrayLike | None = None) -> Array:
  computation_dtype, dtype = _var_promote_types(a.dtype, dtype)
  a = lax.asarray(a).astype(computation_dtype)
  if a_mean is None:
    a_mean = nanmean(a, axis, dtype=computation_dtype, keepdims=True, where=where)
  else:
    a_mean = ensure_arraylike("nanvar", a_mean).astype(computation_dtype)

  centered = _where(lax._isnan(a), 0, lax.sub(a, a_mean))  # double-where trick for gradients.
  if dtypes.issubdtype(centered.dtype, np.complexfloating):
    centered = lax.real(lax.mul(centered, lax.conj(centered)))
  else:
    centered = lax.square(centered)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove out= and use the returned array
2. If allocation churn matters, wrap the call in jax.jit so the compiler handles buffer reuse

Example fix

// before
np.nanvar(x, out=buf)
// after
buf = jnp.nanvar(x)

assert out is None
v = jnp.nanvar(a, axis=axis, ddof=ddof, where=where)

When it happens

Trigger: Calling jnp.nanvar(a, out=buf); also reached indirectly if nanstd forwards out-related state — nanstd calls nanvar internally.

Common situations: Ported NumPy variance-with-NaNs code that used out=; kwargs passthrough from a config-driven stats pipeline.

Related errors

• The 'out' argument to jnp.std is not supported.
• The 'out' argument to jnp.ptp is not supported.
• The 'out' argument to jnp.nanmean is not supported.
• The 'out' argument to jnp.nanstd is not supported.
• The 'out' argument to jnp.{name} is not supported

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