diff --git a/python/paddle/fluid/tests/unittests/test_meshgrid_op.py b/python/paddle/fluid/tests/unittests/test_meshgrid_op.py index eea1ca3282c93a40cc9fcf3149329a358cadcf41..3f03d5ed650a8f5cb1491d7dbf2c6ff98eefd43a 100644 --- a/python/paddle/fluid/tests/unittests/test_meshgrid_op.py +++ b/python/paddle/fluid/tests/unittests/test_meshgrid_op.py @@ -79,7 +79,7 @@ class TestMeshgridOp3(unittest.TestCase): out_2 = np.broadcast_to(out_2, [100, 200]) exe = fluid.Executor(place=fluid.CPUPlace()) - grid_x, grid_y = paddle.tensor.meshgrid([x, y]) + grid_x, grid_y = paddle.tensor.meshgrid(x, y) res_1, res_2 = exe.run(fluid.default_main_program(), feed={'x': input_1, 'y': input_2}, @@ -90,21 +90,72 @@ class TestMeshgridOp3(unittest.TestCase): class TestMeshgridOp4(unittest.TestCase): - def test_errors(self): - with program_guard(Program(), Program()): + def test_list_input(self): + x = fluid.data(shape=[100], dtype='int32', name='x') + y = fluid.data(shape=[200], dtype='int32', name='y') - def test_input_type(): - x = fluid.data(shape=[200], dtype='float32', name='x2') - paddle.tensor.meshgrid(x) + input_1 = np.random.randint(0, 100, [100, ]).astype('int32') + input_2 = np.random.randint(0, 100, [200, ]).astype('int32') - self.assertRaises(TypeError, test_input_type) + out_1 = np.reshape(input_1, [100, 1]) + out_1 = np.broadcast_to(out_1, [100, 200]) + out_2 = np.reshape(input_2, [1, 200]) + out_2 = np.broadcast_to(out_2, [100, 200]) + + exe = fluid.Executor(place=fluid.CPUPlace()) + grid_x, grid_y = paddle.tensor.meshgrid([x, y]) + res_1, res_2 = exe.run(fluid.default_main_program(), + feed={'x': input_1, + 'y': input_2}, + fetch_list=[grid_x, grid_y]) + + assert np.array_equal(res_1, out_1) + assert np.array_equal(res_2, out_2) class TestMeshgridOp5(unittest.TestCase): + def test_tuple_input(self): + x = fluid.data(shape=[100], dtype='int32', name='x') + y = fluid.data(shape=[200], dtype='int32', name='y') + + input_1 = np.random.randint(0, 100, [100, ]).astype('int32') + input_2 = np.random.randint(0, 100, [200, ]).astype('int32') + + out_1 = np.reshape(input_1, [100, 1]) + out_1 = np.broadcast_to(out_1, [100, 200]) + out_2 = np.reshape(input_2, [1, 200]) + out_2 = np.broadcast_to(out_2, [100, 200]) + + exe = fluid.Executor(place=fluid.CPUPlace()) + grid_x, grid_y = paddle.tensor.meshgrid((x, y)) + res_1, res_2 = exe.run(fluid.default_main_program(), + feed={'x': input_1, + 'y': input_2}, + fetch_list=[grid_x, grid_y]) + + assert np.array_equal(res_1, out_1) + assert np.array_equal(res_2, out_2) + + +class TestMeshgridOp6(unittest.TestCase): def test_api_with_dygraph(self): input_3 = np.random.randint(0, 100, [100, ]).astype('int32') input_4 = np.random.randint(0, 100, [200, ]).astype('int32') + with fluid.dygraph.guard(): + tensor_3 = fluid.dygraph.to_variable(input_3) + tensor_4 = fluid.dygraph.to_variable(input_4) + res_3, res_4 = paddle.tensor.meshgrid(tensor_3, tensor_4) + + assert np.array_equal(res_3.shape, [100, 200]) + assert np.array_equal(res_4.shape, [100, 200]) + + +class TestMeshgridOp7(unittest.TestCase): + def test_api_with_dygraph_list_input(self): + input_3 = np.random.randint(0, 100, [100, ]).astype('int32') + input_4 = np.random.randint(0, 100, [200, ]).astype('int32') + with fluid.dygraph.guard(): tensor_3 = fluid.dygraph.to_variable(input_3) tensor_4 = fluid.dygraph.to_variable(input_4) @@ -114,5 +165,19 @@ class TestMeshgridOp5(unittest.TestCase): assert np.array_equal(res_4.shape, [100, 200]) +class TestMeshgridOp7(unittest.TestCase): + def test_api_with_dygraph_tuple_input(self): + input_3 = np.random.randint(0, 100, [100, ]).astype('int32') + input_4 = np.random.randint(0, 100, [200, ]).astype('int32') + + with fluid.dygraph.guard(): + tensor_3 = fluid.dygraph.to_variable(input_3) + tensor_4 = fluid.dygraph.to_variable(input_4) + res_3, res_4 = paddle.tensor.meshgrid((tensor_3, tensor_4)) + + assert np.array_equal(res_3.shape, [100, 200]) + assert np.array_equal(res_4.shape, [100, 200]) + + if __name__ == '__main__': unittest.main() diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py index 85f2e69681ee2b007138db2eeb0944e0382ecf32..b90a7201dc48692160b2810ab63eef390ac066f1 100644 --- a/python/paddle/tensor/creation.py +++ b/python/paddle/tensor/creation.py @@ -704,18 +704,19 @@ def triu(input, diagonal=0, name=None): return _tril_triu_op(LayerHelper('triu', **locals())) -def meshgrid(input, name=None): +def meshgrid(*args, **kwargs): """ :alias_main: paddle.meshgrid :alias: paddle.meshgrid,paddle.tensor.meshgrid,paddle.tensor.creation.meshgrid - This op takes a list of N tensors as input, each of which is 1-dimensional + This op takes a list of N tensors as input *args, each of which is 1-dimensional vector, and creates N-dimensional grids. Args: - input(Variable) : tensors (list of tensor): the shapes of input k tensors are (N1,), + *args(Variable|list of Variable) : tensors (tuple(list) of tensor): the shapes of input k tensors are (N1,), (N2,),..., (Nk,). Support data types: ``float64``, ``float32``, ``int32``, ``int64``. - name (str, optional): The default value is None. Normally there is no need for + **kwargs (optional): Currently, we only accept name in **kwargs + The default value is None. Normally there is no need for user to set this property. For more information, please refer to :ref:`api_guide_Name`. Returns: @@ -735,7 +736,7 @@ def meshgrid(input, name=None): input_2 = np.random.randint(0, 100, [200, ]).astype('int32') exe = fluid.Executor(place=fluid.CPUPlace()) - grid_x, grid_y = paddle.tensor.meshgrid([x, y]) + grid_x, grid_y = paddle.tensor.meshgrid(x, y) res_1, res_2 = exe.run(fluid.default_main_program(), feed={'x': input_1, 'y': input_2}, @@ -749,41 +750,45 @@ def meshgrid(input, name=None): #example 2: in dygraph mode import paddle - import paddle.fluid as fluid import numpy as np + + paddle.enable_imperative() input_3 = np.random.randint(0, 100, [100, ]).astype('int32') input_4 = np.random.randint(0, 100, [200, ]).astype('int32') - with fluid.dygraph.guard(): - tensor_3 = fluid.dygraph.to_variable(input_3) - tensor_4 = fluid.dygraph.to_variable(input_4) - grid_x, grid_y = paddle.tensor.meshgrid([tensor_3, tensor_4]) + tensor_3 = paddle.imperative.to_variable(input_3) + tensor_4 = paddle.imperative.to_variable(input_4) + grid_x, grid_y = paddle.tensor.meshgrid(tensor_3, tensor_4) #the shape of grid_x is (100, 200) #the shape of grid_y is (100, 200) """ + if len(args) == 1 and isinstance(args[0], (list, tuple)): + args = args[0] if in_dygraph_mode(): - num = len(input) - out = core.ops.meshgrid(input, num) + num = len(args) + out = core.ops.meshgrid(list(args), num) return out + name = kwargs.get("name", None) helper = LayerHelper('meshgrid', **locals()) - if not isinstance(input, list): - raise TypeError("The type of input in meshgrid should be list.") + if not isinstance(args, (list, tuple)): + raise TypeError("The type of input args in meshgrid should be list.") - for id, input_ in enumerate(input): + for id, input_ in enumerate(args): check_dtype(input_.dtype, 'create data type', ['float16', 'float32', 'float64', 'int32', 'int64'], 'meshgrid') - num = len(input) + num = len(args) out = [ - helper.create_variable_for_type_inference(dtype=input[i].dtype) + helper.create_variable_for_type_inference(dtype=args[i].dtype) for i in range(num) ] - helper.append_op(type='meshgrid', inputs={'X': input}, outputs={'Out': out}) + helper.append_op( + type='meshgrid', inputs={'X': list(args)}, outputs={'Out': out}) return out