Softmax 算子 API 描述【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench1. 算子简介沿指定维度计算 Softmax 归一化。主要应用场景分类模型的输出层将 logits 转换为概率分布注意力机制中计算注意力权重强化学习中的策略概率输出算子特征难度等级L2Normalization单输入单输出涉及指数运算、求和、除法等多步计算输出元素值在 [0, 1] 范围内沿指定维度求和为 12. 算子定义数学公式基本公式$$ y_i \frac{\exp(x_i)}{\sum_{j}\exp(x_j)} $$数值稳定版本内部实现$$ y_i \frac{\exp(x_i - \max(x))}{\sum_{j}\exp(x_j - \max(x))} $$其中x_i为输入张量沿指定 dim 维度上的第 i 个元素输出满足0 y_i 1且sum(y) 1沿 dim 维度数值稳定版本减去最大值以避免指数溢出3. 接口规范算子原型cann_bench.softmax(Tensor x, int dim) - Tensor y输入参数说明参数类型默认值描述xTensor必选输入张量dimint-1计算 Softmax 的维度输出参数Shapedtype描述y与输入 x 相同与输入 x 相同Softmax 归一化后的张量数据类型x dtype输出 dtypefloat16float16float32float32bfloat16bfloat16规则与约束x 可以为任意维度的张量dim 指定计算 Softmax 的维度支持负数索引输出 shape 和 dtype 与输入完全一致需注意数值稳定性内部实现应使用减最大值技巧避免指数溢出当输入包含 inf 时对应输出为 0 或 1当输入包含 nan 时输出为 nan4. 精度要求采用生态算子精度标准进行验证。误差指标平均相对误差MERE采样点中相对误差平均值$$ \text{MERE} \text{avg}(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$最大相对误差MARE采样点中相对误差最大值$$ \text{MARE} \max(\frac{\text{abs}(actual - golden)}{\text{abs}(golden)\text{1e-7}}) $$通过标准数据类型FLOAT16BFLOAT16FLOAT32HiFLOAT32FLOAT8 E4M3FLOAT8 E5M2通过阈值(Threshold)2^-102^-72^-132^-112^-32^-2当平均相对误差 MERE Threshold最大相对误差 MARE 10 * Threshold 时判定为通过。5. 标准 Golden 代码import torch Softmax 算子 Torch Golden 参考实现 沿指定维度计算 Softmax 归一化 公式: y_i exp(x_i) / sum(exp(x_j)) 参考 PyTorch API: torch.nn.functional.softmax https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html Parameters: - x: 任意维度输入张量 - dim: int, 默认 -1 - 计算 Softmax 的维度 def softmax( x: torch.Tensor, dim: int -1 ) - torch.Tensor: 沿指定维度计算 Softmax 归一化 Args: x: 输入张量任意 shape dim: 计算 Softmax 的维度默认为 -1最后一维 Returns: Softmax 归一化后的张量shape 与输入相同 输出元素值在 [0, 1] 范围内且沿 dim 维度求和为 1 Examples: x torch.randn(1024, 2048) y softmax(x, dim-1) y torch.nn.functional.softmax(x, dimdim) return y6. 额外信息算子调用示例import torch import cann_bench x torch.randn(1024, 2048, dtypetorch.float32, devicenpu) y cann_bench.softmax(x, dim-1) y cann_bench.softmax(x, dim0) y cann_bench.softmax(x, dim1)【免费下载链接】cann-bench评测AI在处理CANN领域代码任务的能力涵盖算子生成、算子优化等领域支撑模型选型、训练效果评估统一量化评估标准识别Agent能力短板构建CANN领域评测平台推动AI能力在CANN领域的持续演进。项目地址: https://gitcode.com/cann/cann-bench创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考