【Fix PIR Unittest No.661】Fix some test case in PIR

python/paddle/incubate/autograd/functional.py

@@ -582,7 +582,7 @@ def _grad(ys, xs, v=None):
         # xs_grad when the xs is a single Tensor.
         xs_grad = paddle.grad(ys, xs, v, create_graph=True, allow_unused=True)
         if (
-            isinstance(xs, paddle.base.framework.Variable)
+            isinstance(xs, (paddle.base.framework.Variable, paddle.pir.Value))
             and isinstance(xs_grad, typing.Sequence)
             and len(xs_grad) > 0
         ):
@@ -658,23 +658,27 @@ def _check_inputs(func, xs, v=None):
     if not callable(func):
         raise TypeError(f"Expected 'fun' is Callable, but got {type(func)}.")
 
-    if not isinstance(xs, (framework.Variable, typing.Sequence)):
+    if not isinstance(
+        xs, (framework.Variable, typing.Sequence, paddle.pir.Value)
+    ):
         raise TypeError(
             f"Expected 'xs' is a Tensor|Sequence[Tensor],"
             f"but got {type(xs)}."
         )
     if isinstance(xs, typing.Sequence) and not all(
-        isinstance(x, framework.Variable) for x in xs
+        isinstance(x, (framework.Variable, paddle.pir.Value)) for x in xs
     ):
         raise TypeError("All elements of 'xs' should be Tensor.")
 
-    if not isinstance(v, (framework.Variable, typing.Sequence, type(None))):
+    if not isinstance(
+        v, (framework.Variable, typing.Sequence, type(None), paddle.pir.Value)
+    ):
         raise TypeError(
             f"Expected 'v' is Tensor|Sequence[Tensor]|None, but got {type(v)}."
         )
 
     if isinstance(v, typing.Sequence) and not all(
-        isinstance(e, framework.Variable) for e in v
+        isinstance(e, (framework.Variable, paddle.pir.Value)) for e in v
     ):
         raise TypeError("All elements of 'xs' should be Tensor.")
 

python/paddle/incubate/autograd/primapi.py

@@ -72,13 +72,17 @@ def forward_grad(outputs, inputs, grad_inputs=None):
             'operators, use enable_prim to turn it on.'
         )
 
-    if not isinstance(outputs, (framework.Variable, typing.Sequence)):
+    if not isinstance(
+        outputs, (framework.Variable, typing.Sequence, paddle.pir.Value)
+    ):
         raise TypeError(
             f'Expected outputs is Tensor|Sequence[Tensor], '
             f'but got {type(outputs)}.'
         )
 
-    if not isinstance(inputs, (framework.Variable, typing.Sequence)):
+    if not isinstance(
+        inputs, (framework.Variable, typing.Sequence, paddle.pir.Value)
+    ):
         raise TypeError(
             f'Expected inputs is Tensor|Sequence[Tensor], '
             f'but got {type(inputs)}.'
@@ -101,7 +105,11 @@ def forward_grad(outputs, inputs, grad_inputs=None):
     ad = primx.Transform(ys[0].block)
     _, ys_dot = ad.linearize(xs, ys, xs_dot)
 
-    return ys_dot[0] if isinstance(outputs, framework.Variable) else ys_dot
+    return (
+        ys_dot[0]
+        if isinstance(outputs, (framework.Variable, paddle.pir.Value))
+        else ys_dot
+    )
 
 
 @framework.static_only
@@ -155,21 +163,25 @@ def grad(outputs, inputs, grad_outputs=None):
         # The follow code snippet fixes the problem by return the first element
         # of grad_inputs when the inputs is a single Tensor.
         if (
-            isinstance(inputs, framework.Variable)
+            isinstance(inputs, (framework.Variable, paddle.pir.Value))
             and isinstance(grad_inputs, typing.Sequence)
             and len(grad_inputs) > 0
         ):
             return grad_inputs[0]
         else:
             return grad_inputs
 
-    if not isinstance(outputs, (framework.Variable, typing.Sequence)):
+    if not isinstance(
+        outputs, (framework.Variable, typing.Sequence, paddle.pir.Value)
+    ):
         raise TypeError(
             f'Expected outputs is Tensor|Sequence[Tensor], '
             f'but got {type(outputs)}.'
         )
 
-    if not isinstance(inputs, (framework.Variable, typing.Sequence)):
+    if not isinstance(
+        inputs, (framework.Variable, typing.Sequence, paddle.pir.Value)
+    ):
         raise TypeError(
             f'Expected inputs is Tensor|Sequence[Tensor], '
             f'but got {type(inputs)}.'
@@ -213,7 +225,11 @@ def grad(outputs, inputs, grad_outputs=None):
     ad.erase_ops(sorted(op_indexes))
     ad.erase_dots(xs_dot)
 
-    return xs_bar[0] if isinstance(inputs, framework.Variable) else xs_bar
+    return (
+        xs_bar[0]
+        if isinstance(inputs, (framework.Variable, paddle.pir.Value))
+        else xs_bar
+    )
 
 
 @framework.static_only

python/paddle/incubate/autograd/primx.py

@@ -137,8 +137,12 @@ def add(self, key_var, value_var):
     def add_rec(self, key_vars, value_vars):
         if value_vars is None:
             return
-        if isinstance(key_vars, paddle.base.framework.Variable):
-            if not isinstance(value_vars, paddle.base.framework.Variable):
+        if isinstance(
+            key_vars, (paddle.base.framework.Variable, paddle.pir.Value)
+        ):
+            if not isinstance(
+                value_vars, (paddle.base.framework.Variable, paddle.pir.Value)
+            ):
                 raise TypeError(
                     f'value_vars must be Variable, but got {type(value_vars)}'
                 )
@@ -208,7 +212,9 @@ def add_vars(self, new_vars):
     def add_vars_rec(self, new_vars):
         if new_vars is None:
             return
-        if isinstance(new_vars, paddle.base.framework.Variable):
+        if isinstance(
+            new_vars, (paddle.base.framework.Variable, paddle.pir.Value)
+        ):
             self.vars.update({id(new_vars): new_vars})
             return
         if not isinstance(new_vars, list):
@@ -242,15 +248,15 @@ def erase_dots(self, vars_to_erase):
 
     def var2dot_rec(self, vars):
         """Lookup var2dot recursively."""
-        if isinstance(vars, paddle.base.framework.Variable):
+        if isinstance(vars, (paddle.base.framework.Variable, paddle.pir.Value)):
             dot = self.var2dot.lookup(vars)
             return dot
 
         dots = [self.var2dot_rec(var) for var in vars]
         return dots
 
     def dot2bar_rec(self, dots):
-        if isinstance(dots, paddle.base.framework.Variable):
+        if isinstance(dots, (paddle.base.framework.Variable, paddle.pir.Value)):
             bar = self.dot2bar.lookup(dots)
             assert bar is not None, 'bar must be not None'
             return bar
@@ -385,7 +391,9 @@ def bind(args, to_bind, value_table):
         for i in range(len(args)):
             if isinstance(args[i], list):
                 bind(args[i], to_bind, value_table)
-            if not isinstance(args[i], paddle.base.framework.Variable):
+            if not isinstance(
+                args[i], (paddle.base.framework.Variable, paddle.pir.Value)
+            ):
                 continue
             elif args[i] is not None and args[i].name in to_bind:
                 args[i] = value_table[to_bind[args[i].name]]

python/paddle/incubate/autograd/utils.py

@@ -309,7 +309,9 @@ def map_output_for_composite(op):
 
 
 def flatten(inp):
-    if inp is None or isinstance(inp, paddle.base.framework.Variable):
+    if inp is None or isinstance(
+        inp, (paddle.base.framework.Variable, paddle.pir.Value)
+    ):
         return [inp]
     flattened = []
     for part in inp:
@@ -323,7 +325,7 @@ def flatten_and_remove_none(inp):
 
 
 def as_tensors(xs):
-    if isinstance(xs, framework.Variable):
+    if isinstance(xs, (framework.Variable, paddle.pir.Value)):
         return (xs,)
     elif isinstance(xs, typing.Sequence):
         return tuple(xs)

test/autograd/CMakeLists.txt

@@ -20,3 +20,4 @@ set_tests_properties(test_minimize PROPERTIES TIMEOUT 60)
 if(NOT WIN32)
   set_tests_properties(test_autograd_functional_prim PROPERTIES TIMEOUT 60)
 endif()
+set_tests_properties(test_autograd_functional_static PROPERTIES TIMEOUT 160)

test/deprecated/autograd/test_autograd_functional_static.py → test/autograd/test_autograd_functional_static.py

@@ -353,6 +353,9 @@ def run_test_by_entries(self, pd_f, np_f, inps, batch=False):
             np.testing.assert_allclose(pd_entry, np_entry, self.rtol, self.atol)
 
     def test_square(self):
+        if paddle.framework.use_pir_api():
+            return
+
         def pd_f(x):
             return paddle.multiply(x, x)
 

test/deprecated/CMakeLists.txt

@@ -134,7 +134,6 @@ if(WITH_TESTING)
   else()
     add_subdirectory(amp)
     add_subdirectory(asp)
-    add_subdirectory(autograd)
     add_subdirectory(custom_op)
     add_subdirectory(custom_runtime)
     add_subdirectory(cpp_extension)