swiglu对外接口

前向接口

输入:

  • x:必选输入,数据类型float16, bfloat16, float32

输出:

  • y:必选输出,数据类型float16, bfloat16, float32

属性:

  • dim:可选属性,数据类型int32_t,缺省-1。

反向接口

输入:

  • dy:必选输入,数据类型float16, bfloat16, float32
  • x:必选输入,数据类型float16, bfloat16, float32

输出:

  • dx:必选输出,数据类型float16, bfloat16, float32

属性:

  • dim:可选属性,数据类型int32_t,缺省-1。

案例

    import os
    import torch
    import torch_npu
    import numpy as np
    import math
    from mindspeed.op_builder import SwigluOpBuilder

    x = np.random.uniform(-2, 2, (8192,1,3904))
    x = torch.from_numpy(x).float().npu()
    y_grad = np.random.uniform(-2, 2, (8192,1,1952))
    y_grad = torch.from_numpy(y_grad).float().npu()

    x.requires_grad = True
    # 正向接口案例
    mindspeed_ops = SwigluOpBuilder().load()
    result = mindspeed_ops.swiglu(x, dim=-1)
    # 反向接口案例
    result.backward(y_grad)