diff --git a/python/tests/ops/test_gather_op.py b/python/tests/ops/test_gather_op.py index 7f200da8d4..519f5bd5f8 100644 --- a/python/tests/ops/test_gather_op.py +++ b/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. # @@ -29,12 +17,14 @@ import unittest import numpy as np from op_test import OpTest, OpTestTool +from op_test_helper import TestCaseHelper import paddle import cinn from cinn.frontend import * 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,82 +34,124 @@ "x86 test will be skipped due to timeout.") class TestGatherOp(OpTest): def setUp(self): - 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 + print(f"\nRunning {self.__class__.__name__}: {self.case}") + self.data = None 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) + inputs = self.case + dtype = self.case["x_dtype"] + 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 = [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 = [out] + self.paddle_outputs.append(out) def build_cinn_program(self, target): + inputs = self.case + dtype = self.case["x_dtype"] + axis = inputs["axis"] 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) - + 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.inputs["x"], self.inputs["index"]], - [out]) + res = self.get_cinn_output(prog, target, [x, index], self.data, [out]) logger.debug(" -- The output of CINN:\n{}".format(res)) - self.cinn_outputs = 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): - 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 +class TestGatherOpAll(TestCaseHelper): + def init_attrs(self): + self.class_name = "TestGatherOpAll" + self.cls = TestGatherOp + # note: The possible values of axis are related to x, so axis is added in self.inputs + 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 = [{ + "x_dtype": "int16", + "y_dtype": "int64" + }, { + "x_dtype": "int32", + "y_dtype": "int64" + }, { + "x_dtype": "int64", + "y_dtype": "int64" + }, { + "x_dtype": "float32", + "y_dtype": "int64" + }, { + "x_dtype": "float64", + "y_dtype": "int64" + }] + self.attrs = [] if __name__ == "__main__": - unittest.main() + TestGatherOpAll().run()