[NJT] Inline through torch.nested.nested_tensor_from_jagged instead o…

…f graph break (#124343) Pull Request resolved: #124343 Approved by: https://github.com/jbschlosser
pytorch · Apr 18, 2024 · ef93402 · ef93402 · pytorchmergebot · Apr 18, 2024
1 parent bbb6e36
commit ef93402
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
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/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,