fix nested

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -231,6 +231,100 @@ def forward(self, x, y):
         self.assertEqualAny(expected, ep.module()(x, y))
         self.assertEqualAny(expected_, ep.module()(-x, y))
 
+    def test_assign_nested_check(self):
+
+        torch_cond = torch.cond
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                def torch_cond_then_3(y, x):
+
+                    def torch_cond_then_1(y, x):
+                        w = x + y
+                        z = x - y
+                        return (w, z)
+
+                    def torch_cond_else_1(y, x):
+                        u = x + 10
+                        w = x + torch.abs(y) + u
+                        z = x - torch.abs(y) + u
+                        return (w, z)
+
+                    w, z = torch_cond(
+                        y.sum() > 0, torch_cond_then_1, torch_cond_else_1, [y, x]
+                    )
+                    return (w, z)
+
+                def torch_cond_else_3(y, x):
+
+                    def torch_cond_then_2(y):
+                        u = y + 1
+                        return u
+
+                    def torch_cond_else_2(y):
+                        u = torch.abs(y) + 10
+                        return u
+
+                    u = torch_cond(y.sum() > 0, torch_cond_then_2, torch_cond_else_2, [y])
+                    w = torch.abs(x) + u
+                    z = torch.abs(x) - u
+                    return (w, z)
+
+                w, z = torch_cond(x.sum() > 0, torch_cond_then_3, torch_cond_else_3, [y, x])
+                return (w, z)
+
+        x, y = torch.rand((3, 4)), torch.rand((3, 4))
+        Model()(x, y)
+
+    def test_rewrite_forward_assign_nested(self):
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                if x.sum() > 0:
+                    if y.sum() > 0:
+                        w = x + y
+                        z = x - y
+                    else:
+                        u = x + 10
+                        w = x + torch.abs(y) + u
+                        z = x - torch.abs(y) + u
+                else:
+                    if y.sum() > 0:
+                        u = y + 1
+                    else:
+                        u = torch.abs(y) + 10
+                    w = torch.abs(x) + u
+                    z = torch.abs(x) - u
+                return w, z
+
+        x, y = torch.rand((3, 4)), torch.rand((3, 4))
+        expected, expected_, expected_0, expected_1 = (
+            Model()(x, y),
+            Model()(-x, y),
+            Model()(x, -y),
+            Model()(-x, -y),
+        )
+
+        rewritten = transform_method(Model.forward, verbose=self.verbose)
+        self.assertIn("torch.abs(", rewritten.code)
+        self.assertIn("abs", rewritten.dump)
+        code = rewritten.code
+        self.assertIn("branch_cond_else_3", code)
+        Model.forward = rewritten.func
+        self.assertEqualAny(expected, Model()(x, y))
+        self.assertEqualAny(expected_, Model()(-x, y))
+        self.assertEqualAny(expected_0, Model()(x, -y))
+        self.assertEqualAny(expected_1, Model()(-x, -y))
+
+        DYN = torch.export.Dim.DYNAMIC
+        ds = ({0: DYN, 1: DYN}, {0: DYN, 1: DYN})
+        ep = torch.export.export(Model(), (x, y), dynamic_shapes=ds)
+        self.assertIn("cond", [str(getattr(n, "target", "?")) for n in ep.graph.nodes])
+        self.assertEqualAny(expected, ep.module()(x, y))
+        self.assertEqualAny(expected_, ep.module()(-x, y))
+        self.assertEqualAny(expected_0, ep.module()(x, -y))
+        self.assertEqualAny(expected_1, ep.module()(-x, -y))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -3,7 +3,6 @@
 import types
 import textwrap
 from typing import Callable, Dict, List, Set, Optional
-import torch
 
 NODE_TYPES = tuple(
     getattr(ast, k)
@@ -34,22 +33,20 @@ def _settl(node, lineno, level=0):
 
 class RewriteControlFlow(ast.NodeTransformer):
     """
-    The class rewrites tests with function ``torch_cond`` :func:`torch.cond`.
+    The class rewrites tests with function :func:`torch.cond`.
     ``empty_tensor`` is a function returning an empty tensor,
     when a branch returns something the other branch does not.
     """
 
     def __init__(
         self,
-        wrapper_name,
-        empty_tensor: str = "make_empty_tensor",
+        prefix: str = "branch_cond",
         skip_objects: Optional[Dict[str, object]] = None,
         args_names: Optional[Set[str]] = None,
     ):
-        self.wrapper_name = wrapper_name
-        self.empty_tensor = empty_tensor
         self.counter = 0
         self.current_func_args = None
+        self.prefix = prefix
         self.skip_objects = skip_objects or {}
         self.args_names = args_names or set()
         self.local_variables = self.args_names.copy()
@@ -83,7 +80,7 @@ def _find_id(self, exprs: List["ast.Node"]) -> List[str]:
             for n in ast.walk(expr):
                 if (
                     isinstance(n, ast.Name)
-                    and isinstance(n.ctx, ast.Store)
+                    and isinstance(n.ctx, ast.Load)
                     and n.id not in self.skip_objects
                 ):
                     vars.append(n.id)
@@ -97,8 +94,8 @@ def _rewrite_if(
         drop = set()
 
         # extract free variables
-        then_name = f"{self.wrapper_name}_then_{self.counter}"
-        else_name = f"{self.wrapper_name}_else_{self.counter}"
+        then_name = f"{self.prefix}_then_{self.counter}"
+        else_name = f"{self.prefix}_else_{self.counter}"
         then_vars = self._find_id(then_exprs)
         else_vars = self._find_id(else_exprs)
         then_else_vars = set(_ for _ in [*then_vars, *else_vars] if _ in known_local_variables)
@@ -173,8 +170,11 @@ def _rewrite_if(
             [ast.Name(id=v, ctx=ast.Load()) for v in then_else_vars],
             ctx=ast.Load(),
         )
+
         call = ast.Call(
-            func=ast.Name(id=self.wrapper_name, ctx=ast.Load()),
+            func=ast.Attribute(
+                value=ast.Name(id="torch", ctx=ast.Load()), attr="cond", ctx=ast.Load()
+            ),
             args=[
                 test_node,
                 ast.Name(id=then_name, ctx=ast.Load()),
@@ -346,8 +346,7 @@ def inplace_add_parent(tree: "ast.Node"):
 
 def transform_method(
     func: Callable,
-    if_name: str = "torch_cond",
-    empty_tensor: str = "make_empty_tensor",
+    prefix: str = "branch_cond",
     verbose: int = 0,
 ) -> RewrittenMethod:
     """
@@ -367,8 +366,7 @@ def transform_method(
         See method ``_filter_target``.
 
     :param func: method or function to rewrite
-    :param if_name: function calling the test
-    :param empty_tensor: function creating an empty tensor
+    :param prefix: prefix used to create the functions for the branches
     :param verbose: verbosity
     :return: rewritten method
 
@@ -390,7 +388,7 @@ def forward(self, x, y):
         x, y = torch.rand((3, 4)), torch.rand((3, 4))
         expected = Model()(x, y)
 
-        rewritten = transform_method(Model.forward, verbose=10)
+        rewritten = transform_method(Model.forward)
         print("-- code --")
         print(rewritten.code)
 
@@ -423,7 +421,7 @@ def forward(self, x, y):
         x, y = torch.rand((3, 4)), torch.rand((3, 4))
         expected = Model()(x, y)
 
-        rewritten = transform_method(Model.forward, verbose=10)
+        rewritten = transform_method(Model.forward)
         print("-- code --")
         print(rewritten.code)
 
@@ -447,8 +445,7 @@ def forward(self, x, y):
         print(f"[transform_method] -- tree --\n\n{ast.dump(tree, indent=2)}")
     # Apply transformation
     transformer = RewriteControlFlow(
-        if_name,
-        empty_tensor=empty_tensor,
+        prefix=prefix,
         skip_objects=modules,
         args_names=set(sig.parameters),
     )
@@ -482,8 +479,6 @@ def forward(self, x, y):
             ) from e
     namespace: Dict[str, type] = {}
     globs = func.__globals__.copy()
-    globs[if_name] = torch.cond
-    globs[empty_tensor] = lambda: torch.tensor([])
     exec(mod, globs, namespace)
     new_func = namespace.get(func.__name__)
     if not isinstance(new_func, types.FunctionType):