功能描述

算子功能：快速高斯误差线性单元激活函数（Fast Gaussian Error Linear Units activation function），对输入的每个元素计算FastGelu；输入是具有任何有效形状的张量。

计算公式
 $\text{[math]}$

接口原型

npu_fast_gelu(Tensor input) -> Tensor

参数说明

input：Tensor类型，即公式中的x。数据类型支持FLOAT16、FLOAT32、BFLOAT16，数据格式支持ND，支持非连续的Tensor。输入最大支持8维。

输出说明

一个Tensor类型的输出，代表fast_gelu的计算结果。

约束说明

input输入中不能含有空指针。
数据类型BFLOAT16仅如下产品型号支持
 Atlas A2训练系列产品/Atlas 800I A2推理产品

。

支持的PyTorch版本

PyTorch 2.2
PyTorch 2.1
PyTorch 1.11

支持的型号

Atlas A2训练系列产品/Atlas 800I A2推理产品

调用示例

单算子调用

import os
import torch
import torch_npu
import numpy as np
data_var = np.random.uniform(0, 1, [4, 2048, 16, 128]).astype(np.float32)
x = torch.from_numpy(data_var).to(torch.float32).npu()
y = torch_npu.npu_fast_gelu(x).cpu().numpy()

图模式调用（图模式目前仅支持PyTorch 2.1版本）

import os
import torch
import torch_npu
import numpy as np
import torch.nn as nn
import torchair as tng
from torchair.configs.compiler_config import CompilerConfig

os.environ["ENABLE_ACLNN"] = "false"
torch_npu.npu.set_compile_mode(jit_compile=True)
class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
    def forward(self, x): 
        y = torch_npu.npu_fast_gelu(x)
        return y
        
npu_mode = Network()
config = CompilerConfig()
npu_backend = tng.get_npu_backend(compiler_config=config)
npu_mode = torch.compile(npu_mode, fullgraph=True, backend=npu_backend, dynamic=False)
data_var = np.random.uniform(0, 1, [4, 2048, 16, 128]).astype(np.float32)
x = torch.from_numpy(data_var).to(torch.float32)
y =npu_mode(x).cpu().numpy()