jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.mean does not support the numpy 'out' in-place output parameter because JAX arrays are immutable. out=None is accepted purely for API compatibility.

Source

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

    if axis is None:
      count = core.dimension_as_value(np.size(a))
    else:
      count = core.dimension_as_value(_axis_size(a, axis))
    count = lax.convert_element_type(count, dtype)
  else:
    count = sum(_broadcast_to(where, np.shape(a)), axis, dtype=dtype, keepdims=keepdims)
  return count

@api.jit(static_argnames=('axis', 'dtype', 'keepdims', 'upcast_f16_for_computation'),
         inline=True)
def _mean(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
          out: None = None, keepdims: bool = False, *,
          upcast_f16_for_computation: bool = True,
          where: ArrayLike | None = None) -> Array:
  a = ensure_arraylike("mean", a)
  where = check_where("mean", where)
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.mean is not supported.")

  if dtype is None:
    result_dtype = dtypes.to_inexact_dtype(a.dtype)
  else:
    result_dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "mean")

  if upcast_f16_for_computation and dtypes.issubdtype(result_dtype, np.inexact):
    computation_dtype = _upcast_f16(result_dtype)
  else:
    computation_dtype = result_dtype

  normalizer = _count(
      a,
      axis=axis,
      keepdims=keepdims,
      where=where,
      dtype=computation_dtype,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the return value: m = jnp.mean(x, axis=0)
2. For buffer updates under jit: buf = buf.at[i].set(jnp.mean(x))

Example fix

// before
jnp.mean(x, axis=0, out=means)
// after
means = jnp.mean(x, axis=0)

means = jnp.mean(x, axis=0)  # never pass out

When it happens

Trigger: jnp.mean(x, out=buf), np.mean(jax_array, out=...) (numpy delegates to the array's .mean method), or x.mean(out=buf) on a JAX array.

Common situations: Accumulator-style loops written for numpy that write means into preallocated arrays; generic stat-computation helpers forwarding kwargs.

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/39a67d1ca64068a2. Report an issue: GitHub.