From 866c28773d2a5c1a85b85f3f34e0d005b41c8970 Mon Sep 17 00:00:00 2001
From: Wang Xinyu <shaywxy@gmail.com>
Date: Tue, 28 Mar 2023 13:03:32 +0800
Subject: [PATCH] [AMP OP&Test] Add float16 OpTest for squeeze, unsqueeze 
 (#52018)

* add squeeze, unsqueeze, transpose fp16 unitest

* Update test_transpose_op.py
---
 .../fluid/tests/unittests/test_squeeze_op.py  | 25 +++++++++++++++++++
 .../tests/unittests/test_unsqueeze_op.py      | 23 +++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/python/paddle/fluid/tests/unittests/test_squeeze_op.py b/python/paddle/fluid/tests/unittests/test_squeeze_op.py
index caedf1430ec..f0400f24667 100755
--- a/python/paddle/fluid/tests/unittests/test_squeeze_op.py
+++ b/python/paddle/fluid/tests/unittests/test_squeeze_op.py
@@ -52,6 +52,31 @@ class TestSqueezeOp(OpTest):
         self.attrs = {"axes": self.axes}
 
 
+class TestSqueezeFP16Op(OpTest):
+    def setUp(self):
+        self.op_type = "squeeze"
+        self.init_test_case()
+        self.inputs = {"X": np.random.random(self.ori_shape).astype("float16")}
+        self.init_attrs()
+        self.outputs = {
+            "Out": self.inputs["X"].reshape(self.new_shape),
+        }
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(["X"], "Out")
+
+    def init_test_case(self):
+        self.ori_shape = (1, 3, 1, 40)
+        self.axes = (0, 2)
+        self.new_shape = (3, 40)
+
+    def init_attrs(self):
+        self.attrs = {"axes": self.axes}
+
+
 class TestSqueezeBF16Op(OpTest):
     def setUp(self):
         self.op_type = "squeeze"
diff --git a/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py b/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py
index 80e6fee5ded..85d21f56464 100755
--- a/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py
+++ b/python/paddle/fluid/tests/unittests/test_unsqueeze_op.py
@@ -50,6 +50,29 @@ class TestUnsqueezeOp(OpTest):
         self.attrs = {"axes": self.axes}
 
 
+class TestUnsqueezeFP16Op(OpTest):
+    def setUp(self):
+        self.init_test_case()
+        self.op_type = "unsqueeze"
+        self.inputs = {"X": np.random.random(self.ori_shape).astype("float16")}
+        self.init_attrs()
+        self.outputs = {"Out": self.inputs["X"].reshape(self.new_shape)}
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(["X"], "Out")
+
+    def init_test_case(self):
+        self.ori_shape = (3, 40)
+        self.axes = (1, 2)
+        self.new_shape = (3, 1, 1, 40)
+
+    def init_attrs(self):
+        self.attrs = {"axes": self.axes}
+
+
 class TestUnsqueezeBF16Op(OpTest):
     def setUp(self):
         self.init_test_case()
-- 
GitLab