jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.logaddexp.reduce (exposed as jnp.logsumexp) does not accept the numpy-style 'out' parameter; JAX arrays are immutable so in-place output buffers are impossible.

Source

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

                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
                    initial=initial, where_=where)


@api.jit(static_argnames=('axis', 'keepdims', 'dtype'), inline=True)
def _reduce_logical_xor(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None,
                        out: None = None, keepdims: bool = False,
                        initial: ArrayLike | None = None, where: ArrayLike | None = None) -> Array:
  return _reduction(a, name="reduce_logical_xor", op=lax.bitwise_xor, init_val=False, preproc=_cast_to_bool,
                    axis=_ensure_optional_axes(axis), dtype=dtype, out=out, keepdims=keepdims,
                    initial=initial, where_=where)


def _logsumexp(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 log(sum(exp(a))) while avoiding precision loss."""
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.logaddexp.reduce is not supported.")
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "jnp.logaddexp.reduce")
  # TODO(phawkins): dtype isn't used here. That seems like a bug!
  del dtype
  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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove out= and assign the return value: res = jnp.logsumexp(x, axis=1)
2. Use buffer.at[idx].set(jnp.logsumexp(...)) to update buffers under jit
3. Strip out from kwargs before forwarding to jnp

Example fix

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

kwargs.pop('out', None)
res = jnp.logsumexp(x, axis=1, **kwargs)

When it happens

Trigger: Calling jnp.logsumexp(x, out=buf) or np.logaddexp.reduce(jax_array, out=...) which delegates to JAX's implementation.

Common situations: Porting scipy/numpy code that reused buffers for logsumexp; generic reduction wrappers that forward **kwargs including out.

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.logaddexp2.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/7b927c8bfa608236. Report an issue: GitHub.