From 5e4f01f5efa6120730480b722ccc6977c5c39d0e Mon Sep 17 00:00:00 2001
From: Li Fuchen <lifuchen@baidu.com>
Date: Mon, 28 Sep 2020 17:23:10 +0800
Subject: [PATCH] modified sample code of add_position_encoding to  2.0
 (#27561)

* modified sample code of add_position_encoding to  2.0, test=document_fix

* use core.op in add_position_encoding API.

* add test for add_position_encoding in dygraph mode
---
 python/paddle/fluid/layers/nn.py              | 18 ++++---
 .../test_add_position_encoding_op.py          | 50 +++++++++++++------
 2 files changed, 46 insertions(+), 22 deletions(-)

diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
index 3f99af1d37c..b5148c215e9 100755
--- a/python/paddle/fluid/layers/nn.py
+++ b/python/paddle/fluid/layers/nn.py
@@ -13177,16 +13177,20 @@ def add_position_encoding(input, alpha, beta, name=None):
     Examples:
         .. code-block:: python
 
-          import paddle.fluid as fluid
+          import numpy as np
+          import paddle
+          import paddle.nn.functional as F
 
-          tensor = fluid.data(
-              name='tensor',
-              shape=[None, 64, 512],
-              dtype='float32')
-          position_tensor = fluid.layers.add_position_encoding(
-              input=tensor, alpha=1.0, beta=1.0)
+          tensor = np.random.randn(16, 32, 64) 
+          tensor = paddle.to_tensor(tensor)
+          position_tensor = F.add_position_encoding(
+                input=tensor, alpha=1.0, beta=1.0)
 
     """
+    if in_dygraph_mode():
+        return core.ops.add_position_encoding(input, "alpha", alpha, "beta",
+                                              beta)
+
     helper = LayerHelper('add_position_encoding', **locals())
     check_variable_and_dtype(input, 'input', ['float32', 'float64'],
                              "add_position_encoding")
diff --git a/python/paddle/fluid/tests/unittests/test_add_position_encoding_op.py b/python/paddle/fluid/tests/unittests/test_add_position_encoding_op.py
index 908b458cb0a..4dc1ed99df6 100644
--- a/python/paddle/fluid/tests/unittests/test_add_position_encoding_op.py
+++ b/python/paddle/fluid/tests/unittests/test_add_position_encoding_op.py
@@ -17,9 +17,31 @@ import math
 import paddle.fluid.core as core
 from op_test import OpTest
 import paddle.fluid as fluid
+import paddle
+import paddle.nn.functional as F
 from paddle.fluid import Program, program_guard
 
 
+def add_position_encoding(input, alpha=1.0, beta=1.0):
+    batch_size = input.shape[0]
+    max_length = input.shape[1]
+    enc_size = input.shape[2]
+
+    out = np.copy(input)
+
+    half_shape = int(enc_size / 2)
+    for i in range(batch_size):
+        for j in range(max_length):
+            for k in range(half_shape):
+                val = j / pow(10000.0, k * 1.0 / (
+                    half_shape - 1)) if half_shape > 1 else j / 10000.0
+                out[i, j, k] = \
+                    input[i, j, k] * alpha + math.sin(val) * beta
+                out[i, j, half_shape + k] = \
+                    input[i, j, half_shape + k] * alpha + math.cos(val) * beta
+    return out
+
+
 class TestAddPositionEncodingTensorOp(OpTest):
     """
     This class is to test the AddPositionEncodingOp
@@ -56,22 +78,8 @@ class TestAddPositionEncodingTensorOp(OpTest):
         self.alpha = 0.6
         self.beta = 0.5
         self.x = np.random.uniform(0.1, 1, [2, 15, 4]).astype(self.dtype)
-        self.out = np.copy(self.x)
 
-        batch_size = self.x.shape[0]
-        max_length = self.x.shape[1]
-        enc_size = self.x.shape[2]
-
-        half_shape = int(enc_size / 2)
-        for i in range(batch_size):
-            for j in range(max_length):
-                for k in range(half_shape):
-                    val = j / pow(10000.0, k * 1.0 / (
-                        half_shape - 1)) if half_shape > 1 else j / 10000.0
-                    self.out[i, j, k] = \
-                        self.x[i, j, k] * self.alpha + math.sin(val) * self.beta
-                    self.out[i, j, half_shape + k] = \
-                        self.x[i, j, half_shape + k] * self.alpha + math.cos(val) * self.beta
+        self.out = add_position_encoding(self.x, self.alpha, self.beta)
 
 
 class TestAddPositionEncodingLoDTensorOp(OpTest):
@@ -145,5 +153,17 @@ class TestAddPositionEncodingOpError(unittest.TestCase):
             self.assertRaises(TypeError, test_Variable)
 
 
+class TestAddPositionEncodingOpDygraph(unittest.TestCase):
+    def test_dygraph(self):
+        paddle.disable_static()
+        tensor = np.random.randn(16, 32, 64)
+        position_tensor = F.add_position_encoding(
+            input=paddle.to_tensor(tensor), alpha=1.0, beta=1.0).numpy()
+        paddle.enable_static()
+
+        position_tensor_np = add_position_encoding(tensor, 1.0, 1.0)
+        self.assertTrue(np.allclose(position_tensor, position_tensor_np))
+
+
 if __name__ == '__main__':
     unittest.main()
-- 
GitLab