support none

Author: xadupre

_unittests/ut_torch_export_patches/test_patch_module.py

@@ -325,6 +325,30 @@ def forward(self, x, y):
         self.assertEqualAny(expected_0, ep.module()(x, -y))
         self.assertEqualAny(expected_1, ep.module()(-x, -y))
 
+    def test_rewrite_forward_none(self):
+
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                if x is None:
+                    x = torch.abs(y)
+                return x + y
+
+        x, y = torch.rand((3, 4)), torch.rand((3, 4))
+        expected, expected_ = 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)
+        Model.forward = rewritten.func
+        self.assertEqualAny(expected, Model()(x, y))
+        self.assertEqualAny(expected_, 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.assertEqualAny(expected, ep.module()(x, y))
+        self.assertEqualAny(expected_, ep.module()(-x, y))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -80,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.Load)
+                    # and isinstance(n.ctx, ast.Load)
                     and n.id not in self.skip_objects
                 ):
                     vars.append(n.id)
@@ -267,9 +267,11 @@ def visit_If(self, node):
                     tuple(t for t in tgt.elts if t.id not in dropped), ctx=ast.Store()
                 )
 
+            added = {tgt.id} if isinstance(tgt, ast.Name) else set(t.id for t in tgt.elts)
             assign = ast.Assign(targets=[tgt], value=call)
             ast.copy_location(assign, node)
             ast.fix_missing_locations(assign)
+            self.local_variables = known_local_variables | added
             return [then_def, else_def, assign]
 
         # Case 2: return in both branches, we assume both branches return the same results.