diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py index 0ae3d1a60a02b1c70880d6604b38f7011a4a70f8..45a5b669be1f6d6c336d45f3b2aa98e26c27209e 100755 --- a/python/paddle/fluid/layers/nn.py +++ b/python/paddle/fluid/layers/nn.py @@ -6306,6 +6306,15 @@ def unsqueeze(input, axes, name=None): """ if in_dygraph_mode(): + if isinstance(axes, int): + axes = [axes] + elif isinstance(axes, Variable): + axes = [axes.numpy().item(0)] + elif isinstance(axes, (list, tuple)): + axes = [ + item.numpy().item(0) if isinstance(item, Variable) else item + for item in axes + ] out, _ = core.ops.unsqueeze2(input, 'axes', axes) return out diff --git a/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py b/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py index 9382d53e7fec6ba9e1217f99ba5006b3dfe5c150..6f713172f1b29d0df8eed212ab1b148d00d7d45e 100644 --- a/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py +++ b/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py @@ -134,29 +134,60 @@ class API_TestUnsqueeze3(unittest.TestCase): result1, = exe.run(feed={"data1": input, "data2": input2}, fetch_list=[result_squeeze]) - self.assertTrue(np.allclose(input1, result1)) + self.assertTrue(np.array_equal(input1, result1)) + self.assertEqual(input1.shape, result1.shape) class API_TestDyUnsqueeze(unittest.TestCase): def test_out(self): with fluid.dygraph.guard(): input_1 = np.random.random([5, 1, 10]).astype("int32") - input1 = np.squeeze(input_1, axis=1) + input1 = np.expand_dims(input_1, axis=1) input = fluid.dygraph.to_variable(input_1) output = paddle.unsqueeze(input, axis=[1]) out_np = output.numpy() - self.assertTrue(np.allclose(input1, out_np)) + self.assertTrue(np.array_equal(input1, out_np)) + self.assertEqual(input1.shape, out_np.shape) class API_TestDyUnsqueeze2(unittest.TestCase): def test_out(self): with fluid.dygraph.guard(): - input_1 = np.random.random([5, 1, 10]).astype("int32") - input1 = np.squeeze(input_1, axis=1) - input = fluid.dygraph.to_variable(input_1) + input1 = np.random.random([5, 10]).astype("int32") + out1 = np.expand_dims(input1, axis=1) + input = fluid.dygraph.to_variable(input1) output = paddle.unsqueeze(input, axis=1) out_np = output.numpy() - self.assertTrue(np.allclose(input1, out_np)) + self.assertTrue(np.array_equal(out1, out_np)) + self.assertEqual(out1.shape, out_np.shape) + + +class API_TestDyUnsqueezeAxisTensor(unittest.TestCase): + def test_out(self): + with fluid.dygraph.guard(): + input1 = np.random.random([5, 10]).astype("int32") + out1 = np.expand_dims(input1, axis=1) + input = fluid.dygraph.to_variable(input1) + output = paddle.unsqueeze(input, axis=paddle.to_tensor([1])) + out_np = output.numpy() + self.assertTrue(np.array_equal(out1, out_np)) + self.assertEqual(out1.shape, out_np.shape) + + +class API_TestDyUnsqueezeAxisTensorList(unittest.TestCase): + def test_out(self): + with fluid.dygraph.guard(): + input1 = np.random.random([5, 10]).astype("int32") + # Actually, expand_dims supports tuple since version 1.18.0 + out1 = np.expand_dims(input1, axis=1) + out1 = np.expand_dims(out1, axis=2) + input = fluid.dygraph.to_variable(input1) + output = paddle.unsqueeze( + fluid.dygraph.to_variable(input1), + axis=[paddle.to_tensor([1]), paddle.to_tensor([2])]) + out_np = output.numpy() + self.assertTrue(np.array_equal(out1, out_np)) + self.assertEqual(out1.shape, out_np.shape) if __name__ == "__main__": diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py index b42a423f3993f3b06f39a9bea1ee5f18d8330557..937ea7cf6715db43099b8b374be0b818acfd203b 100644 --- a/python/paddle/tensor/manipulation.py +++ b/python/paddle/tensor/manipulation.py @@ -746,8 +746,6 @@ def unsqueeze(x, axis, name=None): print(out3.shape) # [1, 1, 1, 5, 10] """ - if isinstance(axis, int): - axis = [axis] return layers.unsqueeze(x, axis, name)