Fix immediate output getattr

coremltools/converters/mil/frontend/torch/test/test_passes.py

@@ -18,6 +18,7 @@
     flatten_graph_input_values,
     flatten_graph_output_values,
     transform_inplace_ops,
+    remove_getattr_nodes
 )
 import coremltools as ct
 
@@ -405,3 +406,52 @@ def forward(self, x):
         y_cm = ct_model.predict({'x': x})['y']
 
         assert((y_cm == np.zeros(shape)).all())
+
+
+    @staticmethod
+    def test_remove_getattr_nodes_immediate_output():
+        graph_nodes = [
+            InternalTorchIRNode(
+                inputs=["self"],
+                attr={"name": "const_out2", "value": None},
+                outputs=["const_out2"],
+                kind="getattr",
+            ),
+            InternalTorchIRNode(
+                inputs=["self"],
+                attr={"name": "const_out1", "value": None},
+                outputs=["const_out1"],
+                kind="getattr",
+            ),
+            InternalTorchIRNode(
+                inputs=["const_out1", "const_out2"],
+                attr={"value": None},
+                outputs=["3"],
+                kind="tupleconstruct",
+            ),
+        ]
+        const2 = torch.tensor([5., 6., 7., 8.])
+        const1 = torch.tensor([1., 2., 3., 4.])
+        graph_params = {'const_out2': const2, 
+                        'const_out1': const1}
+        graph_inputs = []
+        graph_outputs = ["const_out1", "const_out2"]
+
+        graph = InternalTorchIRGraph(
+            nodes=graph_nodes,
+            params=graph_params,
+            inputs=graph_inputs,
+            outputs=graph_outputs,
+        )
+
+        for node in graph.nodes:
+            node.parent = graph
+
+        remove_getattr_nodes(graph)
+
+        np.testing.assert_equal(graph.nodes[0].kind, "constant")
+        np.testing.assert_equal(graph.nodes[1].kind, "constant")
+        np.testing.assert_equal(graph.nodes[2].kind, "tupleconstruct")
+        np.testing.assert_allclose(graph.nodes[0].attr["value"], const2)
+        np.testing.assert_allclose(graph.nodes[1].attr["value"], const1)
+

coremltools/converters/mil/frontend/torch/test/test_torch_conversion_api.py

@@ -929,6 +929,32 @@ def forward(
             past_kv_len += 1
 
 
+    @staticmethod
+    def test_immediate_return_getattr_model():
+        class ImmediateReturnGetAttrModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("my_constant_output", torch.tensor([1.0, 2.0, 3.0, 4.0]))
+                self.register_buffer("my_constant_output2", torch.tensor([5.0, 6.0, 7.0, 8.0]))
+
+            def forward(self, x):
+                # x is a dummy input, not used
+                return self.my_constant_output, self.my_constant_output2
+
+        model = ImmediateReturnGetAttrModel()
+        model.eval()
+        dummy_input = torch.zeros(1)  # Dummy input for tracing
+        traced_model = torch.jit.trace(model, example_inputs=(dummy_input,))
+        mlmodel = ct.convert(
+            traced_model,
+            inputs=[ct.TensorType(shape=(1,))],
+            convert_to='mlprogram'
+        )
+        outputs = mlmodel.predict({"x": np.zeros(1)})
+        assert "my_constant_output" in outputs
+        assert "my_constant_output2" in outputs
+
+
 ###############################################################################
 # Note: Stress tests for PyTorch input / output types
 ###############################################################################

coremltools/converters/mil/frontend/torch/torchir_passes.py

@@ -232,9 +232,9 @@ def forward(self, x):
 def remove_getattr_nodes(graph: InternalTorchIRGraph) -> None:
     """
     Remove the getattr nodes in the graph
+    If they are output nodes, convert them to constant nodes
     """
 
-    getattr_nodes = []
     new_nodes = []
 
     for node in graph.nodes:
@@ -243,16 +243,20 @@ def remove_getattr_nodes(graph: InternalTorchIRGraph) -> None:
             remove_getattr_nodes(block)
 
         if node.kind == "getattr":
-            getattr_nodes.append(node)
+            if node.name in graph.outputs:
+                # create and add new constant node
+                new_nodes.append(
+                        InternalTorchIRNode(
+                            inputs=[],
+                            outputs=node.outputs,
+                            kind="constant",
+                            name="internal_immediate_output_attr",
+                            attr={"value": node.parent.params[node.name].detach()}
+                        )
+                    )
         else:
             new_nodes.append(node)
 
-    # check the getattr nodes not in the outputs
-    for node in getattr_nodes:
-        if node.name in graph.outputs:
-            raise RuntimeError("{} should not be in the graph outputs.".format(node.name))
-
-    # remove the getattr nodes
     graph.nodes = new_nodes