Add gather tests

python/tests/ops/test_gather_op.py

@@ -1,16 +1,4 @@
-# Copyright (c) 2022 CINN Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+#!/usr/bin/env python3
 
 # Copyright (c) 2022 CINN Authors. All Rights Reserved.
 #
@@ -35,6 +23,7 @@
 from cinn.common import *
 import logging
 import os
+from itertools import product
 
 logging.basicConfig(level=os.environ.get('LOG_LEVEL', 'INFO').upper())
 logger = logging.getLogger(name="gather")
@@ -44,81 +33,115 @@
                     "x86 test will be skipped due to timeout.")
 class TestGatherOp(OpTest):
     def setUp(self):
+        self.data = []
         self.init_case()
 
     def init_case(self):
-        self.inputs = {
-            "x":
-            np.array([[[1.1, 1.2, 1.3], [2.1, 2.2, 2.3], [3.1, 3.2, 3.3],
-                       [4.1, 4.2, 4.3]],
-                      [[5.1, 5.2, 5.3], [6.1, 6.2, 6.3], [7.1, 7.2, 7.3],
-                       [8.1, 8.2, 8.3]],
-                      [[9.1, 9.2, 9.3], [10.1, 10.2, 10.3], [11.1, 11.2, 11.3],
-                       [12.1, 12.2, 12.3]]]).astype("float32"),
-            "index":
-            np.array([0, 0, 2, 2]).astype("int32")
-        }
-        self.axis = 0
+        self.inputs = [{"x": [3, 4, 3], "index": [4], "axis": 1}]
+        self.dtypes = ["float32"]
 
     def build_paddle_program(self, target):
-        x = paddle.to_tensor(self.inputs["x"], stop_gradient=True)
-        index = paddle.to_tensor(self.inputs["index"], stop_gradient=True)
-        out = paddle.gather(x, index, self.axis)
-        logger.debug(" -- The output of Paddle:\n{}".format(out))
-        self.paddle_outputs = [out]
+        for inputs, dtype in product(self.inputs, self.dtypes):
+            axis = inputs["axis"]
+            x_shape = inputs["x"]
+            index_shape = inputs["index"]
+            # Paddle does not support negative axis values.
+            axis = axis if axis >= 0 else len(x_shape) + axis
+            x = np.random.randn(*x_shape).astype(dtype)
+            index = np.random.randint(0, x_shape[axis],
+                                      index_shape).astype("int32")
+            self.data.append([x, index])
+            x = paddle.to_tensor(x, stop_gradient=False)
+            index = paddle.to_tensor(index, stop_gradient=False)
+            out = paddle.gather(x, index, axis)
+            logger.debug(" -- The output of Paddle:\n{}".format(out))
+            self.paddle_outputs.append(out)
 
     def build_cinn_program(self, target):
-        builder = NetBuilder("gather")
-        x = builder.create_input(Float(32), self.inputs["x"].shape, "x")
-        index = builder.create_input(
-            Int(32), self.inputs["index"].shape, "index")
-        out = builder.gather(x, index, axis=self.axis)
-
-        prog = builder.build()
-        res = self.get_cinn_output(prog, target, [x, index],
-                                   [self.inputs["x"], self.inputs["index"]],
-                                   [out])
-        logger.debug(" -- The output of CINN:\n{}".format(res))
-        self.cinn_outputs = res
+        for i, (inputs, dtype) in enumerate(product(self.inputs, self.dtypes)):
+            axis = inputs["axis"]
+            builder = NetBuilder("gather")
+            x = builder.create_input(
+                self.nptype2cinntype(dtype), inputs["x"], "x")
+            index = builder.create_input(Int(32), inputs["index"], "index")
+            out = builder.gather(x, index, axis=axis)
+            prog = builder.build()
+            res = self.get_cinn_output(prog, target, [x, index], self.data[i],
+                                       [out])
+            logger.debug(" -- The output of CINN:\n{}".format(res))
+            self.cinn_outputs.extend(res)
 
     def test_check_results(self):
         self.check_outputs_and_grads(all_equal=True)
 
 
-class TestGatherOpCase1(TestGatherOp):
-    def init_case(self):
-        self.inputs = {
-            "x": np.random.random([16, 32, 32]).astype("float32"),
-            "index": np.random.randint(0, 16, 64).astype("int32")
-        }
-        self.axis = 0
-
-
-class TestGatherOpCase2(TestGatherOp):
-    def init_case(self):
-        self.inputs = {
-            "x": np.random.random([16, 32, 32]).astype("float32"),
-            "index": np.random.randint(0, 32, 15).astype("int32")
-        }
-        self.axis = 1
-
-
-class TestGatherOpCase3(TestGatherOp):
-    def init_case(self):
-        self.inputs = {
-            "x": np.random.random([16, 16, 32, 32]).astype("float32"),
-            "index": np.random.randint(0, 32, 8).astype("int32")
-        }
-        self.axis = 2
-
-
-class TestGatherOpCase4(TestGatherOp):
+class TestGatherOpAll(TestGatherOp):
     def init_case(self):
-        self.inputs = {
-            "x": np.random.random([17, 29, 31, 13]).astype("float32"),
-            "index": np.random.randint(0, 13, 11).astype("int32")
-        }
-        self.axis = 3
+        self.inputs = [
+            {
+                "x": [128],
+                "index": [64],
+                "axis": 0
+            },
+            {
+                "x": [16, 32],
+                "index": [32],
+                "axis": 0
+            },
+            {
+                "x": [16, 32],
+                "index": [32],
+                "axis": 1
+            },
+            {
+                "x": [8, 16, 32],
+                "index": [16],
+                "axis": -3
+            },
+            {
+                "x": [8, 16, 32],
+                "index": [8],
+                "axis": -2
+            },
+            {
+                "x": [8, 16, 32],
+                "index": [8],
+                "axis": -1
+            },
+            {
+                "x": [8, 16, 32],
+                "index": [4],
+                "axis": 2
+            },
+            {
+                "x": [16, 8, 4, 64],
+                "index": [4],
+                "axis": 2
+            },
+            {
+                "x": [16, 8, 4, 1024],
+                "index": [4],
+                "axis": 2
+            },
+            {
+                "x": [16, 8, 4, 1],
+                "index": [4],
+                "axis": 2
+            },
+            {
+                "x": [1, 1, 1, 1],
+                "index": [4],
+                "axis": 2
+            },
+        ]
+        self.dtypes = [
+            "float32",
+            "float64",
+            "int16",
+            "int32",
+            "int64",
+            # "uint8"  # note: some types is not supported in paddle now.
+        ]
 
 
 if __name__ == "__main__":