jax-ml/jax · error · NotImplementedError

The 'out' argument to jnp.logaddexp2.reduce is not supported

Error message

The 'out' argument to jnp.logaddexp2.reduce is not supported.

What it means

The base-2 logsumexp (jnp.logaddexp2.reduce) rejects the 'out' parameter for the same reason as other JAX reductions: JAX arrays are immutable and in-place output is unsupported.

Source

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

  a = ensure_arraylike("logsumexp", a)
  where = check_where("logsumexp", where)
  a_arr, = promote_dtypes_inexact(a)
  pos_dims, dims = _reduction_dims(a_arr, axis)
  amax = max(a_arr.real, axis=dims, keepdims=keepdims, where=where, initial=-np.inf)
  amax = lax.stop_gradient(lax.select(lax.is_finite(amax), amax, lax.full_like(amax, 0)))
  amax_with_dims = amax if keepdims else lax.expand_dims(amax, pos_dims)
  exp_a = lax.exp(lax.sub(a_arr, amax_with_dims.astype(a_arr.dtype)))
  sumexp = exp_a.sum(axis=dims, keepdims=keepdims, where=where)
  result = lax.add(lax.log(sumexp), amax.astype(sumexp.dtype))
  return result if initial is None else lax_other.logaddexp(initial, result)


def _logsumexp2(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
                out: None = None, keepdims: bool = False,
                initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
  """Compute log2(sum(2 ** a)) via logsumexp."""
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.logaddexp2.reduce is not supported.")
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(
        dtype, "jnp.logaddexp2.reduce")
  a = ensure_arraylike("logsumexp2", a)
  where = check_where("logsumexp2", where)
  ln2 = float(np.log(2))
  if initial is not None:
    initial *= ln2
  return _logsumexp(a * ln2, axis=axis, dtype=dtype, keepdims=keepdims,
                    where=where, initial=initial) / ln2

@export
def amin(a: ArrayLike, axis: Axis = None, out: None = None,
        keepdims: bool = False, initial: ArrayLike | None = None,
        where: ArrayLike | None = None) -> Array:
  """Alias of :func:`jax.numpy.min`."""
  return min(a, axis=axis, out=out, keepdims=keepdims,
             initial=initial, where=where)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop out= and use the returned array
2. Update buffers with .at[...].set(...) under jit

Example fix

// before
jnp.logaddexp2.reduce(x, axis=1, out=buf)
// after
buf = jnp.logaddexp2.reduce(x, axis=1)

res = jnp.logaddexp2.reduce(x, axis=1)  # no out kwarg

When it happens

Trigger: Calling the reduce method of jnp.logaddexp2 with out=buf, or np.logaddexp2.reduce(jax_array, out=...) delegating to JAX.

Common situations: Numerical-stability code ported from numpy that used out buffers; kwargs-forwarding 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.mean is not supported.

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