register_fx_node_ge_converter

函数功能

将自定义算子注册到TorchAir框架中。

函数原型

def register_fx_node_ge_converter(aten_op)-->ctx

参数说明

参数	输入/输出	说明	是否必选
aten_op	输入	待注册的算子名，例如torch.ops.aten.xxx。	是

返回值说明

返回值用作装饰器，无法独立使用。

约束说明

无

调用示例

import torch
import torch_npu
from torchair import register_fx_node_ge_converter
from torchair.ge import Tensor

@register_fx_node_ge_converter(torch.ops.aten.sin.default)
def conveter_aten_sin_default(self: Tensor, meta_outputs: TensorSpec = None):
    """NB: aten::sin(Tensor self) -> Tensor"""
    raise NotImplementedError("torch.ops.aten.sin.default ge_converter is not implemented!")

父主题： 基本类和接口