From a1350744eb01fd900654734bb3bd22acb314eab3 Mon Sep 17 00:00:00 2001
From: Zhang Ting <709968123@qq.com>
Date: Thu, 23 Jul 2020 10:50:55 +0800
Subject: [PATCH] register fp16 kernel, test=develop (#25630)

---
 paddle/fluid/operators/shape_op.cu | 12 +++++++-----
 python/paddle/fluid/layers/nn.py   |  7 ++++---
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/paddle/fluid/operators/shape_op.cu b/paddle/fluid/operators/shape_op.cu
index 2df4ad1339..4b9dca0d40 100644
--- a/paddle/fluid/operators/shape_op.cu
+++ b/paddle/fluid/operators/shape_op.cu
@@ -14,8 +14,10 @@ limitations under the License. */
 
 #include "paddle/fluid/operators/shape_op.h"
 
-REGISTER_OP_CUDA_KERNEL(shape, paddle::operators::ShapeKernel<int>,
-                        paddle::operators::ShapeKernel<int32_t>,
-                        paddle::operators::ShapeKernel<int64_t>,
-                        paddle::operators::ShapeKernel<float>,
-                        paddle::operators::ShapeKernel<double>);
+REGISTER_OP_CUDA_KERNEL(
+    shape, paddle::operators::ShapeKernel<int>,
+    paddle::operators::ShapeKernel<int32_t>,
+    paddle::operators::ShapeKernel<int64_t>,
+    paddle::operators::ShapeKernel<float>,
+    paddle::operators::ShapeKernel<double>,
+    paddle::operators::ShapeKernel<paddle::platform::float16>);
diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
index 29648fcf15..c511865e0b 100644
--- a/python/paddle/fluid/layers/nn.py
+++ b/python/paddle/fluid/layers/nn.py
@@ -11101,7 +11101,7 @@ def shape(input):
                 input.shape = [3, 2]
 
     Args:
-        input (Variable): The input can be N-D Tensor or SelectedRows with data type float32, float64, int32, int64.
+        input (Variable): The input can be N-D Tensor or SelectedRows with data type float16, float32, float64, int32, int64.
                           If input variable is type of SelectedRows, returns the shape of it's inner tensor.
 
     Returns:
@@ -11124,8 +11124,9 @@ def shape(input):
             res = exe.run(fluid.default_main_program(), feed={'x':img}, fetch_list=[output])
             print(res) # [array([  3, 100, 100], dtype=int32)]
     """
-    check_variable_and_dtype(input, 'input',
-                             ['float32', 'float64', 'int32', 'int64'], 'shape')
+    check_variable_and_dtype(
+        input, 'input', ['float16', 'float32', 'float64', 'int32', 'int64'],
+        'shape')
     helper = LayerHelper('shape', **locals())
     out = helper.create_variable_for_type_inference(dtype='int32')
     helper.append_op(
-- 
GitLab