expand tests fft

heat/fft/tests/test_fft.py

@@ -15,6 +15,15 @@ def test_fft(self):
         self.assertEqual(y.shape, x.shape)
         self.assert_array_equal(y, np_y)
 
+        # 1D distributed
+        x = ht.random.randn(6, split=0)
+        y = ht.fft.fft(x)
+        np_y = np.fft.fft(x.numpy())
+        self.assertIsInstance(y, ht.DNDarray)
+        self.assertEqual(y.shape, x.shape)
+        self.assertTrue(y.split == 0)
+        self.assert_array_equal(y, np_y)
+
         # n-D distributed
         x = ht.random.randn(10, 8, 6, dtype=ht.float64, split=0)
         # FFT along last axis
@@ -25,25 +34,32 @@ def test_fft(self):
         self.assertTrue(y.split == 0)
         self.assert_array_equal(y, np_y)
 
-        # FFT along distributed axis
-        y = ht.fft.fft(x, axis=0)
-        np_y = np.fft.fft(x.numpy(), axis=0)
+        # FFT along distributed axis, n not None
+        n = 8
+        y = ht.fft.fft(x, axis=0, n=n)
+        np_y = np.fft.fft(x.numpy(), axis=0, n=n)
         self.assertIsInstance(y, ht.DNDarray)
-        self.assertEqual(y.shape, x.shape)
+        self.assertEqual(y.shape, np_y.shape)
         self.assertTrue(y.split == 0)
         self.assert_array_equal(y, np_y)
 
         # complex input
         x = x + 1j * ht.random.randn(10, 8, 6, dtype=ht.float64, split=0)
         # FFT along last axis (distributed)
         x.resplit_(axis=2)
-        y = ht.fft.fft(x)
-        np_y = np.fft.fft(x.numpy())
+        y = ht.fft.fft(x, n=n)
+        np_y = np.fft.fft(x.numpy(), n=n)
         self.assertIsInstance(y, ht.DNDarray)
-        self.assertEqual(y.shape, x.shape)
+        self.assertEqual(y.shape, np_y.shape)
         self.assertTrue(y.split == 2)
         self.assert_array_equal(y, np_y)
 
+        # exceptions
+        # wrong input type
+        x = np.random.randn(6, 3, 3)
+        with self.assertRaises(TypeError):
+            ht.fft.fft(x)
+
     def test_ifft(self):
         # 1D non-distributed
         x = ht.random.randn(6)