[thunderfx] Handle output node with no example_value (#2429)

Author: kshitij12345

thunder/dynamo/splitter.py

@@ -172,7 +172,7 @@ def is_thunder_supported_partition(node: torch.fx.Node) -> bool:
             for n in graph_module.graph.nodes:
                 if n.op == "output":
                     for n in n.all_input_nodes:
-                        if n.meta["example_value"].grad_fn is None:
+                        if "example_value" not in n.meta or n.meta["example_value"].grad_fn is None:
                             is_differentiable_outputs.append(False)
                         else:
                             is_differentiable_outputs.append(True)

thunder/tests/test_dynamo.py

@@ -1686,3 +1686,15 @@ def forward(self, x):
         actual_output.backward()
         expected_output.backward()
         torch.testing.assert_close(org_m.fc.weight.grad, thunder_m.fc.weight.grad)
+
+
+def test_thunderfx_node_with_no_example_value():
+    def test_fn(x):
+        y = x + 10
+        z = y.tolist()[0]
+        return z + 2
+
+    x = torch.tensor([1, 2, 3, 4, 5])
+    actual = thunderfx(test_fn)(x)
+    expected = test_fn(x)
+    torch.testing.assert_close(actual, expected)