已合并
[Fix] Registe npu_dtype_cast gradient support complex. #17339
liu-jiaweikf创建于 2025年1月15日
[Fix] Registe npu_dtype_cast gradient support complex. #17339
已合并
从refs/pull/17339/head合入到master
共 2 个文件变更+16-3
| @@ -23,6 +23,7 @@ from .utils import NPU_AUTOGRAD_FUNCTION | |||
| 23 | 23 | ||
| 24 | GRADIENT_IMPLEMENTED_FOR_COMPLEX.update(( | 24 | GRADIENT_IMPLEMENTED_FOR_COMPLEX.update(( |
| 25 | "stft", | 25 | "stft", |
| 26 | + "npu_dtype_cast", | ||
| 26 | )) | 27 | )) |
| 27 | 28 | ||
| 28 | NPU_NATIVEFUNCTIONS = {'npu_format_cast', '_npu_format_cast'} | 29 | NPU_NATIVEFUNCTIONS = {'npu_format_cast', '_npu_format_cast'} |
| @@ -4,6 +4,7 @@ import torch | |||
| 4 | import torch_npu | 4 | import torch_npu |
| 5 | from torch_npu.testing.testcase import TestCase, run_tests | 5 | from torch_npu.testing.testcase import TestCase, run_tests |
| 6 | from torch_npu.testing.common_utils import create_common_tensor | 6 | from torch_npu.testing.common_utils import create_common_tensor |
| 7 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 7 | 8 | ||
| 8 | 9 | ||
| 9 | class TestDtypeCast(TestCase): | 10 | class TestDtypeCast(TestCase): |
| @@ -16,7 +17,7 @@ class TestDtypeCast(TestCase): | |||
| 16 | output = torch_npu.npu_dtype_cast(input1, dst_dtype) | 17 | output = torch_npu.npu_dtype_cast(input1, dst_dtype) |
| 17 | return output.cpu().detach() | 18 | return output.cpu().detach() |
| 18 | 19 | ||
| 19 | - def test_npu_dtype_cast(self, device="npu"): | 20 | + def test_npu_dtype_cast(self): |
| 20 | item = [np.float32, 0, (64, 10)] | 21 | item = [np.float32, 0, (64, 10)] |
| 21 | _, npu_input = create_common_tensor(item, -1, 1) | 22 | _, npu_input = create_common_tensor(item, -1, 1) |
| 22 | dst_dtype = torch.float16 | 23 | dst_dtype = torch.float16 |
| @@ -25,14 +26,25 @@ class TestDtypeCast(TestCase): | |||
| 25 | custom_output = self.custom_op_exec(npu_input, dst_dtype) | 26 | custom_output = self.custom_op_exec(npu_input, dst_dtype) |
| 26 | self.assertRtolEqual(supported_output, custom_output) | 27 | self.assertRtolEqual(supported_output, custom_output) |
| 27 | 28 | ||
| 28 | - def test_npu_dtype_cast_double_backward(self, device="npu"): | 29 | + def test_npu_dtype_cast_double_backward(self): |
| 29 | - x = torch.randn(3, 3, requires_grad=True).to(device) | 30 | + x = torch.randn(3, 3, requires_grad=True).to("npu") |
| 30 | y = torch_npu.npu_dtype_cast(x, torch.half) | 31 | y = torch_npu.npu_dtype_cast(x, torch.half) |
| 31 | z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y)) | 32 | z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y)) |
| 32 | self.assertIsNone(z[0].grad_fn) | 33 | self.assertIsNone(z[0].grad_fn) |
| 33 | z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y), create_graph=True) | 34 | z = torch.autograd.grad(outputs=y, inputs=x, grad_outputs=torch.ones_like(y), create_graph=True) |
| 34 | self.assertIsNotNone(z[0].grad_fn) | 35 | self.assertIsNotNone(z[0].grad_fn) |
| 35 | 36 | ||
| 37 | + | ||
| 38 | + def test_npu_dtype_cast_complex(self): | ||
| 39 | + x = torch.empty([2, 3], dtype=torch.complex64, device="npu") | ||
| 40 | + x.requires_grad_() | ||
| 41 | + y = torch_npu.npu_dtype_cast(x, torch.complex128) | ||
| 42 | + grad_fn = str(y.grad_fn) | ||
| 43 | + self.assertTrue("NpuDtypeCastBackward" in grad_fn) | ||
| 44 | + | ||
| 45 | + with self.assertRaisesRegex(RuntimeError, r'grad can be implicitly created'): | ||
| 46 | + y.sum().backward() | ||
| 47 | + | ||
| 36 | 48 | ||
| 37 | if __name__ == "__main__": | 49 | if __name__ == "__main__": |
| 38 | run_tests() | 50 | run_tests() |