已合并
Bugfix: Custom operators support input keyword arguments #13794
wang-guangbin创建于 2024年8月17日
Bugfix: Custom operators support input keyword arguments #13794
已合并
从refs/pull/13794/head合入到master
共 2 个文件变更+10-0
| @@ -183,6 +183,9 @@ class RegisterCustomSchema: | |||
| 183 | if tags == "{}": | 183 | if tags == "{}": |
| 184 | tag_index = "" | 184 | tag_index = "" |
| 185 | 185 | ||
| 186 | + pattern = r'\bself\b(?=[,\)])' | ||
| 187 | + func_schema = re.sub(pattern, 'input', func_schema) | ||
| 188 | + | ||
| 186 | if f.has_composite_explicit_autograd_kernel: | 189 | if f.has_composite_explicit_autograd_kernel: |
| 187 | name = DispatcherSignature.from_schema(f.func, prefix=f'wrapper_{f.func.name.overload_name}_').name() | 190 | name = DispatcherSignature.from_schema(f.func, prefix=f'wrapper_{f.func.name.overload_name}_').name() |
| 188 | return f'{maybe_tags}m.def({cpp_string(func_schema)}, TORCH_FN(at_npu::native::{name}){tag_index});\n' | 191 | return f'{maybe_tags}m.def({cpp_string(func_schema)}, TORCH_FN(at_npu::native::{name}){tag_index});\n' |
| @@ -29,6 +29,13 @@ class TestFastGelu(TestCase): | |||
| 29 | custom_output = self.custom_op_exec(npu_input) | 29 | custom_output = self.custom_op_exec(npu_input) |
| 30 | self.assertRtolEqual(supported_output, custom_output) | 30 | self.assertRtolEqual(supported_output, custom_output) |
| 31 | 31 | ||
| 32 | + def test_fast_gelu_input_arg(self): | ||
| 33 | + item = [np.float32, 0, [3, 16, 32]] | ||
| 34 | + _, npu_input = create_common_tensor(item, 0, 100) | ||
| 35 | + supported_output = self.supported_op_exec(npu_input) | ||
| 36 | + custom_output = torch_npu.fast_gelu(input=npu_input) | ||
| 37 | + self.assertRtolEqual(supported_output, custom_output) | ||
| 38 | + | ||
| 32 | 39 | ||
| 33 | if __name__ == "__main__": | 40 | if __name__ == "__main__": |
| 34 | run_tests() | 41 | run_tests() |