From 495a9ceb955e596a752e754954cb30f811728403 Mon Sep 17 00:00:00 2001
From: wangchaochaohu <wangchao66@baidu.com>
Date: Tue, 27 Oct 2020 20:14:18 +0800
Subject: [PATCH] fix the input error of size Op (#28272)

---
 python/paddle/fluid/layers/nn.py              |  6 +--
 .../fluid/tests/unittests/test_numel_op.py    |  2 +-
 .../fluid/tests/unittests/test_size_op.py     | 52 +++++++++++++++++++
 3 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
index ac762944b3a..d5157abf1a9 100755
--- a/python/paddle/fluid/layers/nn.py
+++ b/python/paddle/fluid/layers/nn.py
@@ -11363,10 +11363,10 @@ def size(input):
     """
 
     if in_dygraph_mode():
-        return core.ops.size(x)
+        return core.ops.size(input)
     check_variable_and_dtype(
-        x, 'x', ['bool', 'float16', 'float32', 'float64', 'int32', 'int64'],
-        "size")
+        input, 'input',
+        ['bool', 'float16', 'float32', 'float64', 'int32', 'int64'], "size")
     helper = LayerHelper('size', **locals())
     out = helper.create_variable_for_type_inference(dtype='int64')
     helper.append_op(type='size', inputs={'Input': input}, outputs={'Out': out})
diff --git a/python/paddle/fluid/tests/unittests/test_numel_op.py b/python/paddle/fluid/tests/unittests/test_numel_op.py
index d106484d91e..d3b95097957 100644
--- a/python/paddle/fluid/tests/unittests/test_numel_op.py
+++ b/python/paddle/fluid/tests/unittests/test_numel_op.py
@@ -48,7 +48,7 @@ class TestNumelOp2(TestNumelOp):
         self.shape = (0, )
 
 
-class TestNumelOoAPI(unittest.TestCase):
+class TestNumelAPI(unittest.TestCase):
     def test_numel_static(self):
         main_program = fluid.Program()
         startup_program = fluid.Program()
diff --git a/python/paddle/fluid/tests/unittests/test_size_op.py b/python/paddle/fluid/tests/unittests/test_size_op.py
index aec63caa005..09cd35391ba 100644
--- a/python/paddle/fluid/tests/unittests/test_size_op.py
+++ b/python/paddle/fluid/tests/unittests/test_size_op.py
@@ -14,6 +14,8 @@
 
 import unittest
 import numpy as np
+import paddle
+import paddle.fluid as fluid
 from op_test import OpTest
 
 
@@ -53,5 +55,55 @@ class TestLargeTensor(TestSizeOp):
         self.shape = [2**10]
 
 
+class TestSizeAPI(unittest.TestCase):
+    def test_size_static(self):
+        main_program = fluid.Program()
+        startup_program = fluid.Program()
+        with fluid.program_guard(main_program, startup_program):
+            shape1 = [2, 1, 4, 5]
+            shape2 = [1, 4, 5]
+            x_1 = paddle.fluid.data(shape=shape1, dtype='int32', name='x_1')
+            x_2 = paddle.fluid.data(shape=shape2, dtype='int32', name='x_2')
+            input_1 = np.random.random(shape1).astype("int32")
+            input_2 = np.random.random(shape2).astype("int32")
+            out_1 = paddle.fluid.layers.size(x_1)
+            out_2 = paddle.fluid.layers.size(x_2)
+            exe = paddle.static.Executor(place=paddle.CPUPlace())
+            res_1, res_2 = exe.run(feed={
+                "x_1": input_1,
+                "x_2": input_2,
+            },
+                                   fetch_list=[out_1, out_2])
+            assert (np.array_equal(
+                res_1, np.array([np.size(input_1)]).astype("int64")))
+            assert (np.array_equal(
+                res_2, np.array([np.size(input_2)]).astype("int64")))
+
+    def test_size_imperative(self):
+        paddle.disable_static(paddle.CPUPlace())
+        input_1 = np.random.random([2, 1, 4, 5]).astype("int32")
+        input_2 = np.random.random([1, 4, 5]).astype("int32")
+        x_1 = paddle.to_tensor(input_1)
+        x_2 = paddle.to_tensor(input_2)
+        out_1 = paddle.fluid.layers.size(x_1)
+        out_2 = paddle.fluid.layers.size(x_2)
+        assert (np.array_equal(out_1.numpy().item(0), np.size(input_1)))
+        assert (np.array_equal(out_2.numpy().item(0), np.size(input_2)))
+        paddle.enable_static()
+
+    def test_error(self):
+        main_program = fluid.Program()
+        startup_program = fluid.Program()
+        with fluid.program_guard(main_program, startup_program):
+
+            def test_x_type():
+                shape = [1, 4, 5]
+                input_1 = np.random.random(shape).astype("int32")
+                out_1 = paddle.fluid.layers.size(input_1)
+
+            self.assertRaises(TypeError, test_x_type)
+
+
 if __name__ == '__main__':
+    paddle.enable_static()
     unittest.main()
-- 
GitLab