Migrate ReshapeTransform to pir (#60341)

0x45f · web-flow · commit 4af8ecca447e · 2023-12-26T15:34:54.000+08:00
diff --git a/python/paddle/distribution/transform.py b/python/paddle/distribution/transform.py
@@ -166,7 +166,9 @@ def forward(self, x):
         Returns:
             Tensor: Outcome of forward transformation.
         """
-        if not isinstance(x, paddle.base.framework.Variable):
+        if not isinstance(
+            x, (paddle.base.framework.Variable, paddle.pir.Value)
+        ):
             raise TypeError(
                 f"Expected 'x' is a Tensor or Real, but got {type(x)}."
             )
@@ -187,7 +189,9 @@ def inverse(self, y):
         Returns:
             Tensor: Outcome of inverse transform.
         """
-        if not isinstance(y, paddle.base.framework.Variable):
+        if not isinstance(
+            y, (paddle.base.framework.Variable, paddle.pir.Value)
+        ):
             raise TypeError(
                 f"Expected 'y' is a Tensor or Real, but got {type(y)}."
             )
@@ -209,12 +213,14 @@ def forward_log_det_jacobian(self, x):
         Returns:
             Tensor: The log of the absolute value of Jacobian determinant.
         """
-        if not isinstance(x, paddle.base.framework.Variable):
+        if not isinstance(
+            x, (paddle.base.framework.Variable, paddle.pir.Value)
+        ):
             raise TypeError(
                 f"Expected 'y' is a Tensor or Real, but got {type(x)}."
             )
         if (
-            isinstance(x, paddle.base.framework.Variable)
+            isinstance(x, (paddle.base.framework.Variable, paddle.pir.Value))
             and x.dim() < self._domain.event_rank
         ):
             raise ValueError(
@@ -241,7 +247,9 @@ def inverse_log_det_jacobian(self, y):
         Returns:
             Tensor: The value of :math:`log|det J_{f^{-1}}(y)|`.
         """
-        if not isinstance(y, paddle.base.framework.Variable):
+        if not isinstance(
+            y, (paddle.base.framework.Variable, paddle.pir.Value)
+        ):
             raise TypeError(f"Expected 'y' is a Tensor, but got {type(y)}.")
         if y.dim() < self._codomain.event_rank:
             raise ValueError(
diff --git a/test/distribution/test_distribution_transform_static.py b/test/distribution/test_distribution_transform_static.py
@@ -20,6 +20,7 @@
 
 import paddle
 from paddle.distribution import transform, variable
+from paddle.pir_utils import test_with_pir_api
 
 np.random.seed(2022)
 paddle.seed(2022)
@@ -1157,6 +1158,7 @@ def test_domain(self):
     def test_codomain(self):
         self.assertTrue(isinstance(self._t._codomain, variable.Independent))
 
+    @test_with_pir_api
     def test_forward(self):
         exe = paddle.static.Executor()
         sp = paddle.static.Program()
@@ -1177,6 +1179,7 @@ def test_forward(self):
             atol=config.ATOL.get(str(expected.dtype)),
         )
 
+    @test_with_pir_api
     def test_inverse(self):
         exe = paddle.static.Executor()
         sp = paddle.static.Program()
@@ -1198,6 +1201,7 @@ def test_inverse(self):
             atol=config.ATOL.get(str(expected.dtype)),
         )
 
+    @test_with_pir_api
     def test_forward_log_det_jacobian(self):
         exe = paddle.static.Executor()
         sp = paddle.static.Program()
