jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.var does not support the numpy-style 'out' in-place output argument because JAX arrays are immutable; out is accepted only as None for API compatibility with numpy delegation.

Source

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

     [6.22]]
  """
  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("var", a)
  return _var(a, axis=_ensure_optional_axes(axis), dtype=dtype, out=out, correction=correction, keepdims=keepdims,
              where=where, a_mean=mean)

@api.jit(static_argnames=('axis', 'dtype', 'keepdims'))
def _var(a: Array, *, axis: Axis = None, dtype: DTypeLike | None = None,
         out: None = None, correction: int | float = 0, keepdims: bool = False,
         where: ArrayLike | None = None, a_mean: ArrayLike | None = None) -> Array:
  where = check_where("var", where)
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "var")
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.var is not supported.")

  computation_dtype, dtype = _var_promote_types(a.dtype, dtype)
  a = lax.asarray(a).astype(computation_dtype)
  if a_mean is None:
    a_mean = mean(a, axis, dtype=computation_dtype, keepdims=True, where=where)
  else:
    a_mean = ensure_arraylike("var", a_mean).astype(computation_dtype)

  centered = lax.sub(a, a_mean)
  if dtypes.issubdtype(computation_dtype, np.complexfloating):
    centered = lax.real(lax.mul(centered, lax.conj(centered)))
    computation_dtype = centered.dtype  # avoid casting to complex below.
  else:
    centered = lax.square(centered)

  normalizer = _count(
      a,
      axis=axis,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the returned value: v = jnp.var(x, axis=0)
2. Update buffers with .at[...].set(jnp.var(...)) under jit
3. Strip out from forwarded kwargs

Example fix

// before
jnp.var(x, axis=0, out=var_buf)
// after
var_buf = jnp.var(x, axis=0)

v = jnp.var(x, axis=0)  # no out kwarg

When it happens

Trigger: jnp.var(x, out=buf), np.var(jax_array, out=...) via method delegation, or x.var(out=buf) on a JAX array.

Common situations: Statistics pipelines ported from numpy that write into preallocated buffers; kwargs-forwarding variance wrappers.

Related errors

• The 'out' argument to jnp.round is not supported.
• The 'out' argument to jnp.compress is not supported.
• The 'out' argument to jnp.{name} is not supported.
• The 'out' argument to jnp.logaddexp.reduce is not supported.
• The 'out' argument to jnp.logaddexp2.reduce is not supported

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