jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.nanmean does not support the NumPy out= in-place output argument; JAX arrays are immutable so writing into a caller-supplied buffer is impossible. The check fires before any computation.

Source

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

    ...                    [1, 1, 0, 1]], dtype=bool)
    >>> jnp.nanmean(x, axis=1, keepdims=True, where=where)
    Array([[ 3. ],
           [ 9. ],
           [-1.5]], dtype=float32)

    If ``where`` is ``False`` at all elements, ``jnp.nanmean`` returns ``nan``
    along the given axis.

    >>> where = jnp.array([[False],
    ...                    [False],
    ...                    [False]])
    >>> jnp.nanmean(x, axis=0, keepdims=True, where=where)
    Array([[nan, nan, nan, nan]], dtype=float32)
  """
  a = ensure_arraylike("nanmean", a)
  where = check_where("nanmean", where)
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.nanmean is not supported.")
  if dtypes.issubdtype(a.dtype, np.bool_) or dtypes.issubdtype(a.dtype, np.integer):
    return mean(a, axis, dtype, out, keepdims, where=where)
  if dtype is None:
    dtype = dtypes.to_inexact_dtype(a.dtype)
  else:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "mean")
  nan_mask = lax.bitwise_not(lax._isnan(a))
  normalizer = sum(nan_mask, axis=axis, dtype=dtype, keepdims=keepdims, where=where)
  td = lax.div(nansum(a, axis, dtype=dtype, keepdims=keepdims, where=where), normalizer)
  return td


@export
@api.jit(static_argnames=('axis', 'dtype', 'keepdims'))
def nanvar(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None, out: None = None,
           ddof: int = 0, keepdims: bool = False,
           where: ArrayLike | None = None, mean: ArrayLike | None = None) -> Array:
  r"""Compute the variance of array elements along a given axis, ignoring NaNs.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop out and capture the return value
2. Rewrite buffer-reuse idioms as jitted functions; JIT already optimizes allocations

Example fix

// before
np.nanmean(x, out=result)
// after
result = jnp.nanmean(x)

if out is not None:
    out = None  # or raise early with your own message
m = jnp.nanmean(a, axis=axis, dtype=dtype, where=where)

When it happens

Trigger: Calling jnp.nanmean(a, out=buf) or via positional args jnp.nanmean(a, 0, None, True, buf-like).

Common situations: Porting NaN-ignoring statistics code from NumPy/SciPy that used out= for memory reuse; kwargs-forwarding wrappers.

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.nanvar 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/bcab73e32f859cb4. Report an issue: GitHub.