Update on "Improve torch.fft n-dimensional transforms"

[ghstack-poisoned]
pytorch · Dec 8, 2020 · e42964b · e42964b
2 parents 5b6f748 + 3e9686f
commit e42964b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/aten/src/ATen/native/SpectralOps.cpp b/aten/src/ATen/native/SpectralOps.cpp
@@ -301,7 +301,7 @@ Tensor fft_ifftn(const Tensor& self, c10::optional<IntArrayRef> s,
 Tensor fft_rfftn(const Tensor& self, c10::optional<IntArrayRef> s,
                 c10::optional<IntArrayRef> dim,
                 c10::optional<std::string> norm_str) {
-  TORCH_CHECK(!self.is_complex(), "rfftn expects a real input tensor, but got ", self.scalar_type());
+  TORCH_CHECK(!self.is_complex(), "rfftn expects a real-valued input tensor, but got ", self.scalar_type());
   auto desc = canonicalize_fft_shape_and_dim_args(self, s, dim);
   TORCH_CHECK(desc.shape.size() > 0, "rfftn must transform at least one axis");
   Tensor input = promote_tensor_fft(self, /*require_complex=*/false);

diff --git a/test/test_spectral_ops.py b/test/test_spectral_ops.py
@@ -228,7 +228,7 @@ def test_fft_invalid_dtypes(self, device):
         with self.assertRaisesRegex(RuntimeError, "Expected a real input tensor"):
             torch.fft.rfft(t)
 
-        with self.assertRaisesRegex(RuntimeError, "rfftn expects a real input tensor"):
+        with self.assertRaisesRegex(RuntimeError, "rfftn expects a real-valued input tensor"):
             torch.fft.rfftn(t)
 
         with self.assertRaisesRegex(RuntimeError, "Expected a real input tensor"):
@@ -479,7 +479,7 @@ def test_fftn_invalid(self, device):
                 func(a, s=(10, 10, 10, 10))
 
         c = torch.complex(a, a)
-        with self.assertRaisesRegex(RuntimeError, "rfftn expects a real input"):
+        with self.assertRaisesRegex(RuntimeError, "rfftn expects a real-valued input"):
             torch.fft.rfftn(c)
 
     # 2d-fft tests
@@ -591,7 +591,7 @@ def test_fft2_invalid(self, device):
                 func(a, dim=(2, 3))
 
         c = torch.complex(a, a)
-        with self.assertRaisesRegex(RuntimeError, "rfftn expects a real input"):
+        with self.assertRaisesRegex(RuntimeError, "rfftn expects a real-valued input"):
             torch.fft.rfft2(c)
 
     # Helper functions