Fix

Author: co63oc

test/legacy_test/test_fft_op_0size.py

@@ -50,32 +50,21 @@ def test_zero_size(self):
 
 
 op_list = [
-    "fft2",
-    "fftn",
-    "fftshift",
-    "ifft2",
-    "ifftn",
-    "ifftshift",
-    "ihfft2",
-    "ihfftn",
-    "rfft2",
-    "rfftn",
+    ("fft2", ("complex64", "complex128"), ((-2, -1), None)),
+    ("fftn", ("complex64", "complex128"), ((0, 1, 2), None)),
+    ("fftshift", ("complex64", "complex128"), (0, -1)),
+    ("ifft2", ("complex64", "complex128"), ((-2, -1), None)),
+    ("ifftn", ("complex64", "complex128"), ((0, 1, 2), None)),
+    ("ifftshift", ("complex64", "complex128"), (0, -1)),
+    ("ihfft2", ("float32", "float64"), ((-2, -1), None)),
+    ("ihfftn", ("float32", "float64"), ((0, 1, 2), None)),
+    ("rfft2", ("float32", "float64"), ((-2, -1), None)),
+    ("rfftn", ("float32", "float64"), ((0, 1, 2), None)),
 ]
-axes_list = [
-    ((-2, -1), None),
-    (None, None),
-    (0, -1),
-    ((-2, -1), None),
-    (None, None),
-    (0, -1),
-    ((-2, -1), None),
-    (None, None),
-    ((-2, -1), None),
-    (None, None),
-]
-for i, op_type in enumerate(op_list):
-    create_test_class(op_type, "complex64", [3, 4, 0], axes_list[i][0])
-    create_test_class(op_type, "complex128", [3, 4, 0, 3, 4], axes_list[i][1])
+
+for op in op_list:
+    create_test_class(op[0], op[1][0], [3, 4, 0], op[2][0])
+    create_test_class(op[0], op[1][1], [3, 4, 0, 3, 4], op[2][1])
 
 if __name__ == '__main__':
     unittest.main()