From f5997e807bd1b6a9aef42b1a4dc7c4d326a98abd Mon Sep 17 00:00:00 2001 From: Shabalin Denis swx1272397 Date: Sat, 3 Feb 2024 02:06:22 +0800 Subject: [PATCH] Added RFFT test --- test/test_fake_tensor.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/test/test_fake_tensor.py b/test/test_fake_tensor.py index a67833ff80..6c9dbf9ab9 100644 --- a/test/test_fake_tensor.py +++ b/test/test_fake_tensor.py @@ -1493,6 +1493,15 @@ class TestAntiQuant(TestCase): x = x.to(dstType) self.assertTrue(x.numel() * x.element_size() == res.numel() * res.element_size()) +class TestRFFT(TestCase): + def test_npu_rfft(self): + with FakeTensorMode(): + shape = [64,64,1024] + length = shape[-1] + x = torch.randn(shape, dtype=torch.float32).npu() + res = torch.fft.rfft(x, length, norm = "backward") + self.assertTrue(res.shape[2] == (length / 2 + 1)) + instantiate_parametrized_tests(FakeTensorTest) instantiate_device_type_tests(FakeTensorOpInfoTest, globals(), only_for="cpu") -- Gitee