[Dy2St][PIR] Allow define a local variable in single branch (#61936)

SigureMo · web-flow · commit 15daecc16e24 · 2024-02-22T14:22:32.000+08:00
* [Dy2St][PIR] Allow define a local variable in single branch

* cleanup comments
diff --git a/python/paddle/static/nn/control_flow.py b/python/paddle/static/nn/control_flow.py
@@ -1373,6 +1373,7 @@ def restore_outputs_by_variable_results(self, variable_results):
     @staticmethod
     def constant_to_variable_promotion(out_with_blocks, name):
         from paddle.jit.dy2static.convert_operators import to_static_variable
+        from paddle.jit.dy2static.utils import UndefinedVar
 
         promotion_builtin_types = (bool, int, float)
         outs, _ = zip(*out_with_blocks)
@@ -1419,14 +1420,21 @@ def constant_to_variable_with_block(constant, block_context_manager):
         ):
             warnings.warn(
                 "Return results from different branches in cond are not same type: "
-                f"false_var returned by false_fn is '{type(outs[1])}' and true_var of true_fn is "
-                f"'{type(outs[0])}'"
+                + f"false_var returned by false_fn is '{type(outs[1])}' and true_var of true_fn is "
+                + f"'{type(outs[0])}'"
             )
             return [
                 constant_to_variable_with_block(out, block)
                 for out, block in out_with_blocks
             ]
 
+        if any(isinstance(out, UndefinedVar) for out in outs):
+            warnings.warn(
+                f"Return results has maybe unbound local variable `{name}`, please ensure do not use `{name}`"
+                + "after cond."
+            )
+            return [UndefinedVar(name) for _ in out_with_blocks]
+
         raise TypeError(
             "Unsupported return type of true_fn and false_fn in cond: false_var "
             f"returned `{name}` by false_fn is `{outs[0]}` and true_var of true_fn is `{outs[1]}`"
diff --git a/test/dygraph_to_static/test_ifelse.py b/test/dygraph_to_static/test_ifelse.py
@@ -24,6 +24,7 @@
     test_ast_only,
     test_legacy_and_pt_and_pir,
     test_legacy_only,
+    test_pir_only,
 )
 from ifelse_simple_func import (
     NetWithControlFlowIf,
@@ -582,5 +583,54 @@ def test_run_backward(self):
         )
 
 
+def ifelse_temp_local_var(x):
+    if x:
+        y = x + 1
+    else:
+        tmp = x + 2
+        y = tmp * 2
+    return y
+
+
+def ifelse_use_undefined_var(x):
+    if x:
+        y = x + 1
+    else:
+        tmp = x + 2
+        y = tmp * 2
+    return tmp + 1
+
+
+class TestIfElseMaybeUnbound(Dy2StTestBase):
+    @test_legacy_and_pt_and_pir
+    def test_maybe_unbound(self):
+        truethy = paddle.to_tensor(1)
+        falsy = paddle.to_tensor(0)
+
+        dygraph_out = ifelse_temp_local_var(truethy)
+        static_fn = paddle.jit.to_static(ifelse_temp_local_var)
+        static_out = static_fn(truethy)
+        np.testing.assert_allclose(dygraph_out.numpy(), static_out.numpy())
+
+        dygraph_out = ifelse_temp_local_var(falsy)
+        static_fn = paddle.jit.to_static(ifelse_temp_local_var)
+        static_out = static_fn(falsy)
+        np.testing.assert_allclose(dygraph_out.numpy(), static_out.numpy())
+
+    @test_ast_only
+    @test_pir_only
+    def test_use_undefined_var(self):
+        truethy = paddle.to_tensor(1)
+        falsy = paddle.to_tensor(0)
+
+        static_fn = paddle.jit.to_static(ifelse_use_undefined_var)
+        with self.assertRaises(TypeError):
+            static_fn(truethy)
+
+        static_fn = paddle.jit.to_static(ifelse_use_undefined_var)
+        with self.assertRaises(TypeError):
+            static_fn(falsy)
+
+
 if __name__ == '__main__':
     unittest.main()
