[NJT] Inline through torch.nested.nested_tensor_from_jagged instead of graph break (#124343)

soulitzer · pytorchmergebot · commit cf5ca58e7f31 · 2024-04-19T23:13:59.000Z
Pull Request resolved: #124343 Approved by: https://github.com/jbschlosser
diff --git a/test/dynamo/test_subclasses.py b/test/dynamo/test_subclasses.py
@@ -1361,6 +1361,14 @@ def fn(x):
         self._check_recompiles(fn, (nt,), (nt2,), False)
         self._check_recompiles(fn, (nt,), (nt3,), True)
 
+    def test_inline_nested_tensor_from_jagged(self):
+        nt, _ = self._get_jagged_tensor(((2, 3, 4), 5), None)
+
+        def fn(x):
+            return torch.nested.nested_tensor_from_jagged(x.values() * 2, x.offsets())
+
+        torch.compile(fn, fullgraph=True, backend="aot_eager")(nt)
+
     def _get_views(self):
         # Test all cases with both an NT base and a dense base
         # Subclass -> Subclass
diff --git a/test/profiler/test_profiler.py b/test/profiler/test_profiler.py
@@ -634,6 +634,7 @@ def test_execution_trace_no_capture(self):
                 found_root_node = True
         assert found_root_node
 
+    @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/124500")
     def test_execution_trace_nested_tensor(self):
         fp = tempfile.NamedTemporaryFile("w+t", suffix=".et.json", delete=False)
         fp.close()
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
@@ -173,6 +173,7 @@
     "torch.nn.Parameter": TorchInGraphFunctionVariable,
     "torch._nested_tensor_from_mask": SkipFunctionVariable,
     "torch._nested_from_padded": SkipFunctionVariable,
+    "torch.nested.nested_tensor_from_jagged": UserFunctionVariable,
     # symbol operators implemented in Python
     "torch.sym_not": TorchInGraphFunctionVariable,
     "torch.sym_float": TorchInGraphFunctionVariable,
