jax.nn.logsumexp

jax.nn.logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None, keepdims: bool = False, return_sign: Literal[False] = False, where: ArrayLike | None = None) → Array

jax.nn.logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None, keepdims: bool = False, *, return_sign: Literal[True], where: ArrayLike | None = None) → tuple[Array, Array]

jax.nn.logsumexp(a: ArrayLike, axis: Axis = None, b: ArrayLike | None = None, keepdims: bool = False, return_sign: bool = False, where: ArrayLike | None = None) → Array | tuple[Array, Array]

对数-和-指数缩减。

JAX 实现的 scipy.special.logsumexp()。

\[\mathrm{logsumexp}(a) = \mathrm{log} \sum_j b \cdot \mathrm{exp}(a_{ij})\]

其中 \(j\) 索引范围覆盖一个或多个要减少的维度。

参数:

• a – 输入数组

• axis – 要减少的轴或轴。可以是 None、一个整数或一个整数的元组。

• b – 用于 \(\mathrm{exp}(a)\) 的缩放因子。必须可广播到 a 的形状。

• keepdims – 如果 True ，被减少的轴将作为大小为1的维度保留在输出中。

• return_sign – 如果 True，输出将是一个 (result, sign) 对，其中 sign 是和的符号，而 result 包含它们绝对值的对数。如果 False，则只返回 result，如果和为负数，它将包含 NaN 值。

• where – 要包含在简化中的元素。

返回:

根据 return_sign 参数的值，返回一个数组 result 或一对数组 (result, sign)。