确定算子原型
本样例目标是自定义实现一个非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完全相同。
请先完成环境准备,确定好算子原型后,实现目标算子入图的步骤如下:

父主题: 非In-place算子开发和入图样例