From 5bcabf784dfd374d434a9172732948c756aea98d Mon Sep 17 00:00:00 2001
From: houj04 <35131887+houj04@users.noreply.github.com>
Date: Tue, 16 Aug 2022 11:44:48 +0800
Subject: [PATCH] [XPU] add truncated_gaussian_random op. (#45152)

---
 .../fluid/platform/device/xpu/xpu2_op_list.h  |  2 +
 .../test_truncated_gaussian_random_op_xpu.py  | 97 +++++++++++++++++--
 2 files changed, 90 insertions(+), 9 deletions(-)

diff --git a/paddle/fluid/platform/device/xpu/xpu2_op_list.h b/paddle/fluid/platform/device/xpu/xpu2_op_list.h
index bb20116e234..eb7889a11b5 100644
--- a/paddle/fluid/platform/device/xpu/xpu2_op_list.h
+++ b/paddle/fluid/platform/device/xpu/xpu2_op_list.h
@@ -558,6 +558,8 @@ XPUOpMap& get_kl2_ops() {
       {"transpose",
        XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace()),
                      pOpKernelType(vartype::FP16, XPUPlace())})},
+      {"truncated_gaussian_random",
+       XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace())})},
       {"top_k",
        XPUKernelSet({pOpKernelType(vartype::FP32, XPUPlace()),
                      pOpKernelType(vartype::FP16, XPUPlace())})},
diff --git a/python/paddle/fluid/tests/unittests/xpu/test_truncated_gaussian_random_op_xpu.py b/python/paddle/fluid/tests/unittests/xpu/test_truncated_gaussian_random_op_xpu.py
index f9ccf0576a2..9dd6f8a4f8a 100644
--- a/python/paddle/fluid/tests/unittests/xpu/test_truncated_gaussian_random_op_xpu.py
+++ b/python/paddle/fluid/tests/unittests/xpu/test_truncated_gaussian_random_op_xpu.py
@@ -1,4 +1,4 @@
-#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,24 +18,103 @@ import sys
 
 sys.path.append("..")
 import unittest
-import numpy
-
+import numpy as np
 import paddle
 import paddle.fluid as fluid
-import paddle.fluid.core as core
-from paddle.fluid.op import Operator
 from paddle.fluid.executor import Executor
-from test_truncated_gaussian_random_op import TestTrunctedGaussianRandomOp
+from op_test_xpu import XPUOpTest
+from xpu.get_test_cover_info import create_test_class, get_xpu_op_support_types, XPUOpTestWrapper
 
 paddle.enable_static()
 
 
-class TestXPUTrunctedGaussianRandomOp(TestTrunctedGaussianRandomOp):
+class XPUTestTruncatedGaussianRandomOp(XPUOpTestWrapper):
+
+    def __init__(self):
+        self.op_name = 'truncated_gaussian_random'
+        self.use_dynamic_create_class = False
+
+    class TestTruncatedGaussianRandomOp(XPUOpTest):
+
+        def init(self):
+            self.dtype = self.in_type
+            self.place = paddle.XPUPlace(0)
+            self.__class__.op_type = "truncated_gaussian_random"
+
+        def setUp(self):
+            self.init()
+            self.inputs = {}
+            self.set_attrs()
+            self.attrs = {
+                "shape": self.shape,
+                "mean": self.mean,
+                "std": self.std,
+                "seed": 10,
+            }
+            self.outputs = {'Out': np.zeros(self.shape, dtype=self.dtype)}
 
-    def test_xpu(self):
-        if paddle.is_compiled_with_xpu():
+        def set_attrs(self):
+            self.shape = [10000]
+            self.mean = 0.0
+            self.std = 1.0
+
+        def test_check_output(self):
             self.gaussian_random_test(place=fluid.XPUPlace(0))
 
+        def gaussian_random_test(self, place):
+
+            program = fluid.Program()
+            block = program.global_block()
+            vout = block.create_var(name="Out")
+            op = block.append_op(type=self.op_type,
+                                 outputs={"Out": vout},
+                                 attrs=self.attrs)
+
+            op.desc.infer_var_type(block.desc)
+            op.desc.infer_shape(block.desc)
+
+            fetch_list = []
+            for var_name in self.outputs:
+                fetch_list.append(block.var(var_name))
+
+            exe = Executor(place)
+            outs = exe.run(program, fetch_list=fetch_list)
+            tensor = outs[0]
+            np.testing.assert_allclose(np.mean(tensor), self.mean, atol=0.05)
+            np.testing.assert_allclose(np.var(tensor), 0.773, atol=0.05)
+
+    class TestTruncatedGaussianRandomOp_1(TestTruncatedGaussianRandomOp):
+
+        def set_attrs(self):
+            self.shape = [4096, 2]
+            self.mean = 5.0
+            self.std = 1.0
+
+    class TestTruncatedGaussianRandomOp_2(TestTruncatedGaussianRandomOp):
+
+        def set_attrs(self):
+            self.shape = [1024]
+            self.mean = -2.0
+            self.std = 1.0
+
+    class TestTruncatedGaussianRandomOp_3(TestTruncatedGaussianRandomOp):
+
+        def set_attrs(self):
+            self.shape = [11 * 13 * 17]
+            self.mean = -1.0
+            self.std = 1.0
+
+    class TestTruncatedGaussianRandomOp_4(TestTruncatedGaussianRandomOp):
+
+        def set_attrs(self):
+            self.shape = [2049]
+            self.mean = 5.1234
+            self.std = 1.0
+
+
+support_types = get_xpu_op_support_types('truncated_gaussian_random')
+for stype in support_types:
+    create_test_class(globals(), XPUTestTruncatedGaussianRandomOp, stype)
 
 if __name__ == "__main__":
     unittest.main()
-- 
GitLab