mlx.core.softmax

softmax(a: array, /, axis: None | int | Sequence[int] = None, *, stream: None | Stream | Device = None) → array

沿给定轴执行 softmax 操作。

此操作是以下公式的数值稳定版本：

exp(a) / sum(exp(a), axis, keepdims=True)

参数:

• a (array) – 输入数组。

• axis (int 或 list(int), 可选) – 计算 softmax 的可选轴或多个轴。如果未指定，则对整个数组执行 softmax。

返回值:

softmax 的输出。

返回类型:

array