Added reshape grad bf16 (#31035)

* - added Reshape grad bf16 * - Added reshape grad bf16 * - cosmetics in py
PaddlePaddle · Feb 19, 2021 · f746564 · f746564
1 parent 4dbe16c
commit f746564
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/paddle/fluid/operators/reshape_op.cc b/paddle/fluid/operators/reshape_op.cc
@@ -642,12 +642,14 @@ REGISTER_OP_CPU_KERNEL_FUNCTOR(
     reshape2_grad, float, ops::ReshapeGradKernel, double,
     ops::ReshapeGradKernel, int, ops::ReshapeGradKernel, uint8_t,
     ops::ReshapeGradKernel, int64_t, ops::ReshapeGradKernel, bool,
-    ops::ReshapeGradKernel, paddle::platform::complex64, ops::ReshapeGradKernel,
+    ops::ReshapeGradKernel, paddle::platform::bfloat16, ops::ReshapeGradKernel,
+    paddle::platform::complex64, ops::ReshapeGradKernel,
     paddle::platform::complex128, ops::ReshapeGradKernel);
 REGISTER_OP_CPU_KERNEL_FUNCTOR(
     reshape2_grad_grad, float, ops::ReshapeDoubleGradKernel, double,
     ops::ReshapeDoubleGradKernel, int, ops::ReshapeDoubleGradKernel, uint8_t,
     ops::ReshapeDoubleGradKernel, int64_t, ops::ReshapeDoubleGradKernel, bool,
+    ops::ReshapeDoubleGradKernel, paddle::platform::bfloat16,
     ops::ReshapeDoubleGradKernel, paddle::platform::complex64,
     ops::ReshapeDoubleGradKernel, paddle::platform::complex128,
     ops::ReshapeDoubleGradKernel);

diff --git a/python/paddle/fluid/tests/unittests/mkldnn/test_reshape_bf16_op.py b/python/paddle/fluid/tests/unittests/mkldnn/test_reshape_bf16_op.py
@@ -28,7 +28,7 @@
 class TestReshapeBf16Op(OpTest):
     def setUp(self):
         self.op_type = "reshape2"
-        self.use_mkldnn = True
+        self.use_mkldnn = False
         self.mkldnn_data_type = "bfloat16"
         self.init_data()
         self.init_input_data()
@@ -56,6 +56,16 @@ def init_input_data(self):
     def test_check_output(self):
         self.check_output_with_place(core.CPUPlace(), no_check_set=['XShape'])
 
+    def test_check_grad(self):
+        self.check_grad_with_place(
+            core.CPUPlace(), ["X"],
+            "Out",
+            check_dygraph=False,
+            user_defined_grads=[self.inputs["X"]],
+            user_defined_grad_outputs=[
+                self.inputs["X"].reshape(self.infered_shape)
+            ])
+
 
 if __name__ == '__main__':
     enable_static()