昇腾社区首页
中文
注册
开发者
下载

确定算子原型

本样例目标是自定义实现一个非In-place类PyTorch算子,使其在Eager和TorchAir图模式下正常工作。

假设目标PyTorch算子Schema定义如下:

定义非In-place类算子时,需要遵循PyTorch原型定义惯例:

  • Tensor类型输入在前,基本类型输入在后。
- func: my_op(Tensor x, Tensor? y, Tensor[] z, float attr1, int attr2) -> Tensor
  • my_op:算子名,对应调用方式为torch.ops.npu.my_op。
  • Tensor x:表示x为Tensor类型输入。
  • Tensor? y:表示y为Tensor类型输入,但是允许传None。
  • Tensor[] z:表示z为List[Tensor]类型输入。
  • float attr1:表示attr1为float类型的值输入。
  • int attr2:表示attr2为int类型的值输入。
  • -> Tensor:表示返回一个Tensor。

输出Tensor的shape与dtype与算子逻辑相关,假设本样例的输出Tensor shape和dtype与输入x完全相同。

请先完成环境准备,确定好算子原型后,实现目标算子入图的步骤如下: