[dynamic shapes][export] ignore when real-tensor fallback fails (pytorch#147779)

Summary: uninspired solution to pytorch#147402

Test Plan: test_draft_export

Differential Revision: D70132269

Pull Request resolved: pytorch#147779
Approved by: https://github.com/bobrenjc93

Author: pianpwk

test/export/test_draft_export.py

@@ -386,6 +386,21 @@ def forward(self, x):
         self.assertTrue(report.successful())
         self.assertEqual(inp[0], torch.ones(3, 3))
 
+    def test_masked_linear(self):
+        class M(torch.nn.Module):
+            def forward(self, x, mask, weight, bias):
+                masked = x[mask != 0, :, :]
+                return torch.nn.functional.linear(masked, weight, bias)
+
+        x = torch.zeros(10)
+        inp = (torch.randn(10, 8, 7), x, torch.randn(25, 7), torch.randn(25))
+        draft_ep = draft_export(M(), inp)
+        ep = export(M(), inp)
+        self.assertEqual(draft_ep.module()(*inp), ep.module()(*inp))
+        x[2] += 1
+        x[3] += 1
+        self.assertEqual(draft_ep.module()(*inp), ep.module()(*inp))
+
     def test_torchbind(self):
         class Model(torch.nn.Module):
             def __init__(self) -> None:

torch/_subclasses/fake_tensor.py

@@ -2201,7 +2201,17 @@ def maybe_to_real_tensor(
                     func, real_flat_args, args_spec
                 )
 
-            real_out = func(*real_args, **real_kwargs)
+            try:
+                real_out = func(*real_args, **real_kwargs)
+            except ZeroDivisionError as exc:
+                # we shouldn't broadly catch all errors here;
+                # some come from real-kernel mutation/aliasing checks we want to run.
+                # add more exception types as needed.
+                log.debug(
+                    "real-tensor fallback failed for %s: %s; silently ignoring",
+                    func,
+                    exc,
+                )
 
             if not is_builtin:
                 mutation_checker.check()  # type: ignore[possibly-undefined]