jax.nn.softmax

jax.nn.softmax(x, axis=-1, where=None)

Softmax 函数。

计算将元素重新缩放到范围 \([0, 1]\) 的函数，使沿 axis 的元素之和为 \(1\)。

\[\mathrm{softmax}(x) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}\]

参数:

• x (ArrayLike) – 输入数组

• axis (Axis) – 计算softmax的轴或多个轴。这些维度上的softmax输出总和应为 \(1\)。可以是整数或整数元组。

• where (ArrayLike | None) – 包含在 softmax 中的元素。任何被遮罩的元素的输出为零。

返回:

一个数组。

返回类型:

数组

注意

如果任何输入值为 +inf，结果将全是 NaN：这反映了 inf / inf 在浮点运算中未明确定义的事实。

另请参阅

log_softmax()