【PIR API adaptor No.271】Migrate LogNormal to pir (#60318)

0x45f · web-flow · commit 8710cb794f4e · 2023-12-28T14:31:56.000+08:00
diff --git a/python/paddle/distribution/distribution.py b/python/paddle/distribution/distribution.py
@@ -150,7 +150,7 @@ def _validate_args(self, *args):
         is_variable = False
         is_number = False
         for arg in args:
-            if isinstance(arg, Variable):
+            if isinstance(arg, (Variable, paddle.pir.Value)):
                 is_variable = True
             else:
                 is_number = True
@@ -176,7 +176,10 @@ def _to_tensor(self, *args):
         tmp = 0.0
 
         for arg in args:
-            if not isinstance(arg, (float, list, tuple, np.ndarray, Variable)):
+            if not isinstance(
+                arg,
+                (float, list, tuple, np.ndarray, Variable, paddle.pir.Value),
+            ):
                 raise TypeError(
                     "Type of input args must be float, list, tuple, numpy.ndarray or Tensor, but received type {}".format(
                         type(arg)
diff --git a/python/paddle/distribution/normal.py b/python/paddle/distribution/normal.py
@@ -99,13 +99,29 @@ def __init__(self, loc, scale, name=None):
             check_type(
                 loc,
                 'loc',
-                (int, float, np.ndarray, Variable, list, tuple),
+                (
+                    int,
+                    float,
+                    np.ndarray,
+                    Variable,
+                    paddle.pir.Value,
+                    list,
+                    tuple,
+                ),
                 'Normal',
             )
             check_type(
                 scale,
                 'scale',
-                (int, float, np.ndarray, Variable, list, tuple),
+                (
+                    int,
+                    float,
+                    np.ndarray,
+                    Variable,
+                    paddle.pir.Value,
+                    list,
+                    tuple,
+                ),
                 'Normal',
             )
 
diff --git a/test/distribution/test_distribution_lognormal_static.py b/test/distribution/test_distribution_lognormal_static.py
@@ -33,9 +33,10 @@
         ('one-dim', xrand((2,)), xrand((2,)), xrand((2,))),
         ('multi-dim', xrand((3, 3)), xrand((3, 3)), xrand((3, 3))),
     ],
+    test_pir=True,
 )
 class TestLogNormal(unittest.TestCase):
-    def setUp(self):
+    def run_program(self):
         paddle.enable_static()
         startup_program = paddle.static.Program()
         main_program = paddle.static.Program()
@@ -67,6 +68,13 @@ def setUp(self):
             self.log_prob,
         ] = executor.run(main_program, feed=self.feeds, fetch_list=fetch_list)
 
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.run_program()
+        else:
+            self.run_program()
+
     def test_mean(self):
         np_mean = self.np_lognormal.mean
         self.assertEqual(str(self.mean.dtype).split('.')[-1], self.scale.dtype)
@@ -122,9 +130,10 @@ def test_log_prob(self):
 @parameterize_cls(
     (TEST_CASE_NAME, 'loc', 'scale'),
     [('sample', xrand((4,)), xrand((4,), min=0, max=1))],
+    test_pir=True,
 )
 class TestLogNormalSample(unittest.TestCase):
-    def setUp(self):
+    def run_program(self):
         paddle.enable_static()
         startup_program = paddle.static.Program()
         main_program = paddle.static.Program()
@@ -150,6 +159,13 @@ def setUp(self):
             main_program, feed=self.feeds, fetch_list=fetch_list
         )
 
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.run_program()
+        else:
+            self.run_program()
+
     def test_sample(self):
         samples_mean = self.samples.mean(axis=0)
         samples_var = self.samples.var(axis=0)
@@ -196,9 +212,10 @@ def _kstest(self, loc, scale, samples):
             xrand((2, 2)),
         ),
     ],
+    test_pir=True,
 )
 class TestLogNormalKL(unittest.TestCase):
-    def setUp(self):
+    def run_program(self):
         paddle.enable_static()
         startup_program = paddle.static.Program()
         main_program = paddle.static.Program()
@@ -236,6 +253,13 @@ def setUp(self):
             main_program, feed=self.feeds, fetch_list=fetch_list
         )
 
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.run_program()
+        else:
+            self.run_program()
+
     def test_kl_divergence(self):
         np.testing.assert_allclose(
             self.kl0,
