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

确定算子原型

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

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

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

  • 必须显式标记被In-place修改的Tensor输入,标记方式例如Tensor(a!),a为别名,!表示该输入被修改。
  • Tensor类型输入在前,基本类型输入在后。
  • 不返回任何被In-place修改的输入。

为什么遵循这个规则,因为本样例基于PyTorch 2.1.0版本,需要手动实现函数化转换。

考虑到PyTorch 2.5+版本开始支持函数化自动转换,要求不返回被修改的输入,TorchAir将在后续版本支持该特性。

- func: my_inplace(Tensor(a!) x, Tensor(b!) y) -> ()
  • my_inplace:算子名,对应调用方式为torch.ops.npu.my_inplace。
  • Tensor(a!) x:表示x为Tensor类型输入,(a!)表示该输入具有别名a且会被修改。
  • Tensor(b!) y:表示y为Tensor类型输入,(b!)表示该输入具有别名b且会被修改。
  • -> ():表示算子无返回值,对应Python中的None。

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