Softmax是一种归一化激活函数，可以将Tensor（张量）的某个或多个维度数据归一化，对应维度上每一个元素的范围都在(0, 1)之间并且和为1。

图1 SoftmaxOperation

$\text{[math]}$

通常作为模型的最后一层，例如在Transformer模型中将Self-Attention的K、Q进行softmax，获取相应的相关性。

$\text{[math]}$

      
           struct SoftmaxParam {
    SVector<int64_t> axes;
    uint8_t rsv[8] = {0};
};

成员名称	类型	默认值	描述
axes	SVector<int64_t>	-	指定轴（维度），axes可以支持多个轴上进行处理。 axes不能为空，当指定多个轴时，多个轴之间必须连续且从小到大排列。 axes的元素要求大于或等于-1且小于输入x的维度。
rsv[8]	uint8_t	{0}	预留参数。

成员名称

类型

默认值

描述

axes

SVector<int64_t>

-

指定轴（维度），axes可以支持多个轴上进行处理。

rsv[8]

uint8_t

{0}

预留参数。

参数	维度	数据类型	格式	描述
x	[dim_0, dim_1, ..., dim_n]	float16/float/bf16	ND	输入tensor。

参数	维度	数据类型	格式	描述
output	[dim_0, dim_1, ..., dim_n]	float16/float/bf16	ND	输出tensor。维度和数据类型与输入tensor一致。

单维度softmax
当参数axes只含一个元素时，对输入tensor的对应维度进行softmax操作。当axes=[x]时，等效于Pytorch的

$\text{[math]}$
多维度softmax
当参数axes含有多个元素，且符合要求时，对输入tensor的对应多个维度进行softmax操作。axes=[x, x+1]时，对x和x+1维度对应的元素进行softmax操作。将axes中连续的多个维度平铺成一维，然后对这一维进行softmax操作，操作完成后再reshape回原来的维度形状。