Refactor test_tensordot (#42650)

* Refactor test_tensordot

* Add test_static

* Fix CI errors

python/paddle/fluid/tests/unittests/CMakeLists.txt

@@ -1036,6 +1036,7 @@ set_tests_properties(test_imperative_selected_rows_to_lod_tensor PROPERTIES TIME
 set_tests_properties(test_index_select_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_parallel_ssa_graph_inference_feed_partial_data PROPERTIES TIMEOUT 120)
 set_tests_properties(test_parallel_executor_crf PROPERTIES TIMEOUT 120)
+set_tests_properties(test_tensordot PROPERTIES TIMEOUT 200)
 set_tests_properties(test_imperative_save_load PROPERTIES TIMEOUT 120)
 set_tests_properties(test_partial_eager_deletion_transformer PROPERTIES TIMEOUT 120)
 set_tests_properties(test_parallel_executor_seresnext_with_reduce_gpu PROPERTIES TIMEOUT 120)
@@ -1233,9 +1234,6 @@ if(WITH_GPU OR WITH_ROCM)
 endif()
 set_tests_properties(test_inplace_addto_strategy PROPERTIES TIMEOUT 120)
 set_tests_properties(test_eigvals_op PROPERTIES TIMEOUT 400)
-set_tests_properties(test_tensordot PROPERTIES TIMEOUT 1000)
-set_tests_properties(test_tensordot PROPERTIES LABELS "RUN_TYPE=NIGHTLY")
-set_tests_properties(test_tensordot PROPERTIES ENVIRONMENT "FLAGS_USE_STANDALONE_EXECUTOR=False")
 set_tests_properties(test_cuda_memory_reserved PROPERTIES ENVIRONMENT "FLAGS_allocator_strategy=auto_growth")
 if (WITH_GLOO)
     set_tests_properties(test_parallel_dygraph_dataparallel_cpuonly PROPERTIES TIMEOUT 30)

python/paddle/fluid/tests/unittests/test_tensordot.py

@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import paddle
-import unittest
-import paddle.fluid.core as core
-import numpy as np
 import itertools as it
+import numpy as np
+import unittest
 
-np.set_printoptions(threshold=np.inf)
+import paddle
+import paddle.fluid.core as core
 
 
 def tensordot_np(x, y, axes):
@@ -68,9 +67,16 @@ def tensordot_np(x, y, axes):
 
 class TestTensordotAPI(unittest.TestCase):
     def setUp(self):
+        self.set_place()
         self.set_dtype()
         self.set_input_shape()
         self.set_input_data()
+        self.set_test_axes()
+
+    def set_place(self):
+        self.places = [core.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            self.places.append(core.CUDAPlace(0))
 
     def set_dtype(self):
         self.dtype = np.float32
@@ -82,33 +88,8 @@ class TestTensordotAPI(unittest.TestCase):
     def set_input_data(self):
         self.x = np.random.random(self.x_shape).astype(self.dtype)
         self.y = np.random.random(self.y_shape).astype(self.dtype)
-        self.all_axes = [2]
 
-    def run_dygraph(self, place):
-        paddle.disable_static()
-        x = paddle.to_tensor(self.x, place=place)
-        y = paddle.to_tensor(self.y, place=place)
-        paddle_res = paddle.tensordot(x, y, self.axes)
-        np_res = tensordot_np(self.x, self.y, self.axes)
-        np.testing.assert_allclose(paddle_res, np_res, rtol=1e-6)
-
-    def run_static(self, place):
-        paddle.enable_static()
-        with paddle.static.program_guard(paddle.static.Program(),
-                                         paddle.static.Program()):
-            x = paddle.static.data(
-                name='x', shape=self.x_shape, dtype=self.dtype)
-            y = paddle.static.data(
-                name='y', shape=self.y_shape, dtype=self.dtype)
-            z = paddle.tensordot(x, y, self.axes)
-            exe = paddle.static.Executor(place)
-            paddle_res = exe.run(feed={'x': self.x,
-                                       'y': self.y},
-                                 fetch_list=[z])
-            np_res = tensordot_np(self.x, self.y, self.axes)
-            np.testing.assert_allclose(paddle_res[0], np_res, rtol=1e-6)
-
-    def test_cases(self):
+    def set_test_axes(self):
         self.all_axes = []
         axial_index = range(4)
         all_permutations = list(it.permutations(axial_index, 0)) + list(
@@ -136,57 +117,146 @@ class TestTensordotAPI(unittest.TestCase):
 
         self.all_axes.extend(range(5))
 
-        places = [core.CPUPlace()]
-        if core.is_compiled_with_cuda():
-            places.append(core.CUDAPlace(0))
-
+    def test_dygraph(self):
+        paddle.disable_static()
+        for axes in self.all_axes:
+            for place in self.places:
+                x = paddle.to_tensor(self.x, place=place)
+                y = paddle.to_tensor(self.y, place=place)
+                paddle_res = paddle.tensordot(x, y, axes)
+                np_res = tensordot_np(self.x, self.y, axes)
+                np.testing.assert_allclose(paddle_res, np_res, rtol=1e-6)
+
+    def test_static(self):
+        paddle.enable_static()
         for axes in self.all_axes:
-            self.axes = axes
-            for place in places:
-                self.run_dygraph(place)
-                self.run_static(place)
+            for place in self.places:
+                with paddle.static.program_guard(paddle.static.Program(),
+                                                 paddle.static.Program()):
+                    x = paddle.static.data(
+                        name='x', shape=self.x_shape, dtype=self.dtype)
+                    y = paddle.static.data(
+                        name='y', shape=self.y_shape, dtype=self.dtype)
+                    z = paddle.tensordot(x, y, axes)
+                    exe = paddle.static.Executor(place)
+                    paddle_res = exe.run(feed={'x': self.x,
+                                               'y': self.y},
+                                         fetch_list=[z])
+                    np_res = tensordot_np(self.x, self.y, axes)
+                    np.testing.assert_allclose(paddle_res[0], np_res, rtol=1e-6)
 
 
 class TestTensordotAPIFloat64(TestTensordotAPI):
+    # Only test a small part of axes case for Float64 type
+    def set_test_axes(self):
+        self.all_axes = [
+            [[3, 2], [3]], [[2, 1, 0], [2, 1]], [[1, 2, 0], [1, 3, 2]], [3, 0],
+            [[], [0, 3, 1]], [[2, 1, 0, 3], [2, 0, 1, 3]],
+            [[3, 1, 2], [1, 3, 2, 0]], [[2, 1], [0, 2]], [[2, 0, 1, 3], [2]],
+            [[1, 2, 0, 3], [0, 2, 1]], [[2, 1, 3, 0], [1, 2, 3]],
+            [[2, 0, 1, 3], [3, 1, 0, 2]], [[0, 3], [0, 3, 2, 1]],
+            [[1, 3, 2, 0], [2, 1, 0, 3]], [[1, 3, 2, 0], [1, 3, 2, 0]],
+            [[1, 0, 2], [0, 1]], [[2, 3, 0], [3, 1]],
+            [[1, 3, 2, 0], [3, 0, 1, 2]], [[3, 2, 1], [2, 0, 1]], [[0], []],
+            [[2, 3, 0], [1, 2, 0]], [[3, 0, 2, 1], [2, 1, 0, 3]],
+            [[3, 1, 2], [2, 3, 1]], [[1, 0, 2, 3], []], [[1, 2], [1, 2, 3]],
+            [[2, 0, 1, 3], [2, 0, 1]], [[3, 1, 2], [1, 3, 2]],
+            [[3, 1, 2, 0], [1, 2, 3, 0]], [[0, 2, 3], [0, 1, 2]],
+            [[3, 2, 0], [2, 0, 3, 1]], [[2, 1, 0, 3], [3, 1, 2, 0]],
+            [[1, 2, 3, 0], [1, 3, 0, 2]], [[3, 0], [2, 1]],
+            [[0, 1, 3, 2], [0, 2, 1, 3]], [[1, 0], [2, 1, 3]],
+            [[1, 0, 3, 2], [2, 3, 0, 1]], [[1, 2], [3]],
+            [[1, 2, 3, 0], [3, 2, 1, 0]], [[0, 3, 2, 1], [2, 1, 3, 0]], [0],
+            [[0, 2, 3], [3, 2, 0, 1]], [[1, 2, 3, 0], [3, 2, 1, 0]],
+            [[3, 1], [3]], [[3, 2, 0, 1], [3, 2, 0]], [[2, 3, 0, 1], [0, 3, 2]],
+            [[1], [1, 3]], [[1, 2], [2, 1, 0]], [[3, 1, 2], [3, 1, 0]],
+            [[1, 3], [3, 1, 2]], [[2, 0, 1, 3], [3, 1, 0, 2]],
+            [[1, 3, 0], [1, 3]], [[2, 3, 1], [1, 0, 2]],
+            [[1, 2, 0, 3], [0, 2, 1, 3]], [[2], [0, 1, 3]], [[1], [1, 2]],
+            [[1, 0, 2, 3], [3, 0, 1, 2]], [[0, 1, 3, 2], [1, 3, 0, 2]],
+            [[3, 0, 2, 1], [0, 2, 3]], [[1, 2, 0], [1, 2, 3]],
+            [[1, 0, 3], [2, 3, 0]], [[2, 3, 0], [3, 1, 0]], [[1, 3], [1, 0]],
+            [[2, 1, 0, 3], [2, 0, 3, 1]], [[3, 2, 0], [2, 1, 0]],
+            [[0, 1, 3], [0, 3, 1]], [[3, 1, 0], [3, 2, 1]], [[3, 2], [3, 1]],
+            [[3], [2, 1, 0]], [[1, 2, 3, 0], []], [[1, 3, 2, 0], [3, 1, 2]],
+            [[1], [0, 2]], [[3, 2, 0], [3, 2, 0]], [[3], []],
+            [[1, 0, 3], [2, 1]], [[3, 1, 0, 2], [2, 3, 1, 0]],
+            [[0, 1], [0, 3, 2]], [[0, 2, 3], [0, 2, 1]], [[1, 3, 0], [3, 0, 2]],
+            [[3, 1, 2], [1, 2, 3]], [[3, 1, 2], [3, 1, 0]],
+            [[0, 3, 1, 2], [3, 2, 1, 0]], [[0, 3], [3, 2, 1]],
+            [[2, 3], [1, 3, 0]], [[0, 3, 2], [2, 0, 3, 1]], [[2, 3], [1, 3]],
+            [[3, 1, 2, 0], [2, 3, 1, 0]], [[1, 0, 3, 2], [3, 0, 1, 2]],
+            [[3, 2, 1, 0], [0, 1, 3, 2]], [[3, 1, 2], [3]],
+            [[0, 1, 3, 2], [2, 3, 0, 1]], [[1, 2, 3, 0], [1, 3, 0, 2]],
+            [3, 1, 2], [[3, 1, 2], [0, 3, 2]], [[2, 3, 0], [1, 2, 0]],
+            [[2, 0, 3], [2, 0]], [[3, 1, 0, 2], [3, 1, 0, 2]],
+            [[0, 1, 2], [2, 0, 1]], [[1, 0, 3], [2, 3, 0]],
+            [[2, 0, 1], [0, 1, 3]], [[2, 1], [0, 1, 3]]
+        ]
+
     def set_dtype(self):
         self.dtype = np.float64
 
 
+class TestTensordotAPIBroadcastCase1(TestTensordotAPIFloat64):
+    def set_input_shape(self):
+        self.x_shape = [1, 1, 1, 5]
+        self.y_shape = [1, 5, 1, 1]
+
+
+class TestTensordotAPIBroadcastCase2(TestTensordotAPIFloat64):
+    def set_input_shape(self):
+        self.x_shape = [1, 5, 5, 5]
+        self.y_shape = [1, 1, 1, 5]
+
+
+class TestTensordotAPIBroadcastCase3(TestTensordotAPIFloat64):
+    def set_input_shape(self):
+        self.x_shape = [5, 5, 5, 1]
+        self.y_shape = [5, 5, 1, 5]
+
+
+class TestTensordotAPIBroadcastCase4(TestTensordotAPIFloat64):
+    def set_input_shape(self):
+        self.x_shape = [5, 5, 5, 1]
+        self.y_shape = [1, 1, 1, 1]
+
+
+class TestTensordotAPIBroadcastCase5(TestTensordotAPIFloat64):
+    def set_input_shape(self):
+        self.x_shape = [1, 1, 5, 5]
+        self.y_shape = [5, 5, 1, 5]
+
+
 class TestTensordotAPIAxesType(TestTensordotAPI):
     def set_input_shape(self):
         self.x_shape = [3, 4, 4]
         self.y_shape = [4, 4, 5]
 
-    def test_cases(self):
+    def set_test_axes(self):
         self.all_axes = [
             0, 1, 2, (1, ), [1], ((1, ), ), ([1], ), ((2, 1), (0, )), (
                 (1, 2), (0, 1)), ([1, 2], [0, 1]), ([1, 2], [0, 1]),
             [[1, 2], [0, 1]]
         ]
 
-        places = [core.CPUPlace()]
-        if core.is_compiled_with_cuda():
-            places.append(core.CUDAPlace(0))
-
-        for axes in self.all_axes:
-            self.axes = axes
-            for place in places:
-                self.run_dygraph(place)
-                self.run_static(place)
-
+    def test_tensor_axes(self):
         # The 'axes' with type 'Tensor' in tensordot is not available in static mode
         paddle.disable_static()
-        for place in places:
-            self.all_axes = [
-                paddle.to_tensor([1]), (paddle.to_tensor([1])),
-                (paddle.to_tensor([1, 2]), paddle.to_tensor([0, 1])),
-                [paddle.to_tensor([1, 2]), paddle.to_tensor([0, 1])],
-                paddle.to_tensor([[1, 2], [0, 1]])
-            ]
-            for axes in self.all_axes:
-                self.axes = axes
-                for place in places:
-                    self.run_dygraph(place)
+        tensor_axes = [
+            paddle.to_tensor([1]), (paddle.to_tensor([1])),
+            (paddle.to_tensor([1, 2]), paddle.to_tensor([0, 1])),
+            [paddle.to_tensor([1, 2]), paddle.to_tensor([0, 1])],
+            paddle.to_tensor([[1, 2], [0, 1]])
+        ]
+
+        for place in self.places:
+            for axes in tensor_axes:
+                x = paddle.to_tensor(self.x, place=place)
+                y = paddle.to_tensor(self.y, place=place)
+                paddle_res = paddle.tensordot(x, y, axes)
+                np_res = tensordot_np(self.x, self.y, axes)
+                np.testing.assert_allclose(paddle_res, np_res, rtol=1e-6)
 
     def test_error(self):
         self.all_axes = [[[[0], [1]]], 0.1, -1, 100, [[1, 2], [0, 0]],
@@ -204,35 +274,5 @@ class TestTensordotAPIAxesTypeFloat64(TestTensordotAPIAxesType):
         self.dtype = np.float64
 
 
-class TestTensordotAPIBroadcastCase1(TestTensordotAPI):
-    def set_input_shape(self):
-        self.x_shape = [1, 1, 1, 5]
-        self.y_shape = [1, 5, 1, 1]
-
-
-class TestTensordotAPIBroadcastCase2(TestTensordotAPI):
-    def set_input_shape(self):
-        self.x_shape = [1, 5, 5, 5]
-        self.y_shape = [1, 1, 1, 5]
-
-
-class TestTensordotAPIBroadcastCase3(TestTensordotAPI):
-    def set_input_shape(self):
-        self.x_shape = [5, 5, 5, 1]
-        self.y_shape = [5, 5, 1, 5]
-
-
-class TestTensordotAPIBroadcastCase4(TestTensordotAPI):
-    def set_input_shape(self):
-        self.x_shape = [5, 5, 5, 1]
-        self.y_shape = [1, 1, 1, 1]
-
-
-class TestTensordotAPIBroadcastCase5(TestTensordotAPI):
-    def set_input_shape(self):
-        self.x_shape = [1, 1, 5, 5]
-        self.y_shape = [5, 5, 1, 5]
-
-
 if __name__ == "__main__":
     unittest.main()