tico 0.1.0.dev250616__py3-none-any.whl → 0.1.0.dev250618__py3-none-any.whl

tico/passes/decompose_group_norm.py

@@ -23,6 +23,7 @@ from torch.export import ExportedProgram
 
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -89,24 +90,40 @@ class DecomposeGroupNorm(PassBase):
     def __init__(self):
         super().__init__()
 
-    def _insert_norm(self, graph, tensor, eps):
+    def _insert_norm(self, graph, tensor, eps, origin):
         """
         Insert (tensor - mean) / sqrt(var + eps)) into the graph
           and return the normalized tensor node.
         """
-        mean = graph.call_function(
-            torch.ops.aten.mean.dim, (tensor, [-1]), {"keepdim": True}
+        mean = create_node(
+            graph,
+            torch.ops.aten.mean.dim,
+            (tensor, [-1]),
+            {"keepdim": True},
+            origin=origin,
         )
-        deviation = graph.call_function(torch.ops.aten.sub.Tensor, (tensor, mean))
-        squared = graph.call_function(torch.ops.aten.pow.Tensor_Scalar, (deviation, 2))
-        var = graph.call_function(
-            torch.ops.aten.mean.dim, (squared, [-1]), {"keepdim": True}
+        deviation = create_node(
+            graph, torch.ops.aten.sub.Tensor, (tensor, mean), origin=origin
         )
-        inverse_std = graph.call_function(
+        squared = create_node(
+            graph, torch.ops.aten.pow.Tensor_Scalar, (deviation, 2), origin=origin
+        )
+        var = create_node(
+            graph,
+            torch.ops.aten.mean.dim,
+            (squared, [-1]),
+            {"keepdim": True},
+            origin=origin,
+        )
+        inverse_std = create_node(
+            graph,
             torch.ops.aten.rsqrt.default,
-            (graph.call_function(torch.ops.aten.add.Tensor, (var, eps)),),
+            (create_node(graph, torch.ops.aten.add.Tensor, (var, eps), origin=origin),),
+            origin=origin,
+        )
+        return create_node(
+            graph, torch.ops.aten.mul.Tensor, (deviation, inverse_std), origin=origin
         )
-        return graph.call_function(torch.ops.aten.mul.Tensor, (deviation, inverse_std))
 
     def call(self, exported_program: ExportedProgram) -> PassResult:
         logger = logging.getLogger(__name__)
@@ -178,17 +195,23 @@ class DecomposeGroupNorm(PassBase):
                 # Branch only on whether a reshape is needed; the normalization is shared.
                 if norm_size != x_shape[-1]:
                     # Pack groups so that the last dimension equals norm_size.
-                    packed = graph.call_function(
-                        torch.ops.aten.reshape.default, (x, pack_shape)
+                    packed = create_node(
+                        graph,
+                        torch.ops.aten.reshape.default,
+                        (x, pack_shape),
+                        origin=node,
                     )
-                    normed = self._insert_norm(graph, packed, eps)
+                    normed = self._insert_norm(graph, packed, eps, origin=node)
                     # Restore the original shape after normalization.
-                    layer_norm = graph.call_function(
-                        torch.ops.aten.reshape.default, (normed, x_shape)
+                    layer_norm = create_node(
+                        graph,
+                        torch.ops.aten.reshape.default,
+                        (normed, x_shape),
+                        origin=node,
                     )
                 else:
                     # The input already has norm_size in the last dimension.
-                    layer_norm = self._insert_norm(graph, x, eps)
+                    layer_norm = self._insert_norm(graph, x, eps, origin=node)
 
                 # weight
                 if weight:
@@ -197,13 +220,17 @@ class DecomposeGroupNorm(PassBase):
                         assert weight_shape[0] == C
                         reshape_size = [1] * len(x_shape)
                         reshape_size[1] = C
-                        weight = graph.call_function(
+                        weight = create_node(
+                            graph,
                             torch.ops.aten.view.default,
                             (weight, reshape_size),
+                            origin=node,
                         )
-                    layer_norm = graph.call_function(
+                    layer_norm = create_node(
+                        graph,
                         torch.ops.aten.mul.Tensor,
                         (layer_norm, weight),
+                        origin=node,
                     )
 
                 # bias
@@ -213,15 +240,17 @@ class DecomposeGroupNorm(PassBase):
                         assert bias_shape[0] == C
                         reshape_size = [1] * len(x_shape)
                         reshape_size[1] = C
-                        bias = graph.call_function(
+                        bias = create_node(
+                            graph,
                             torch.ops.aten.view.default,
                             (bias, reshape_size),
+                            origin=node,
                         )
-                    layer_norm = graph.call_function(
+                    layer_norm = create_node(
+                        graph,
                         torch.ops.aten.add.Tensor,
                         (layer_norm, bias),
                     )
-
                 # Reset last node's meta for propagating replacing node's meta.
                 layer_norm.meta = {}
 

tico/passes/decompose_grouped_conv2d.py

@@ -23,7 +23,7 @@ from tico.passes import ops
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.errors import InvalidArgumentError, NotYetSupportedError
-from tico.utils.graph import add_placeholder
+from tico.utils.graph import add_placeholder, create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -159,19 +159,26 @@ class DecomposeGroupedConv2d(PassBase):
 
                 conv2d_tensors = []
                 for i in range(groups):
-                    sliced_input = graph.call_function(
+                    sliced_input = create_node(
+                        graph,
                         torch.ops.aten.slice.Tensor,
                         (input_, 1, group_size * i, group_size * (i + 1), 1),
+                        origin=node,
                     )
-                    sliced_weight = graph.call_function(
+                    sliced_weight = create_node(
+                        graph,
                         torch.ops.aten.slice.Tensor,
                         (weight, 0, out_group_size * i, out_group_size * (i + 1), 1),
+                        origin=node,
                     )
-                    sliced_bias = graph.call_function(
+                    sliced_bias = create_node(
+                        graph,
                         torch.ops.aten.slice.Tensor,
                         (bias, 0, out_group_size * i, out_group_size * (i + 1), 1),
+                        origin=node,
                     )
-                    conv2d_tensor = graph.call_function(
+                    conv2d_tensor = create_node(
+                        graph,
                         conv2d_op,
                         (
                             sliced_input,
@@ -182,11 +189,12 @@ class DecomposeGroupedConv2d(PassBase):
                             dilation,
                             1,
                         ),
+                        origin=node,
                     )
                     conv2d_tensors.append(conv2d_tensor)
 
-                concat_output = graph.call_function(
-                    torch.ops.aten.cat.default, (conv2d_tensors, 1)
+                concat_output = create_node(
+                    graph, torch.ops.aten.cat.default, (conv2d_tensors, 1)
                 )
 
                 node.replace_all_uses_with(concat_output, propagate_meta=True)

tico/passes/decompose_slice_scatter.py

@@ -23,6 +23,7 @@ from torch.export import ExportedProgram
 from tico.serialize.circle_mapping import extract_shape
 
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import enforce_type, is_target_node
@@ -130,16 +131,19 @@ class DecomposeSliceScatter(PassBase):
                 slices = []
 
                 if 0 < start:
-                    slice_0 = graph.call_function(
+                    slice_0 = create_node(
+                        graph,
                         torch.ops.aten.slice_copy.Tensor,
                         args=(input, dim, 0, start, 1),
+                        origin=node,
                     )
                     slices.append(slice_0)
 
                 slices.append(src)
 
                 if start + end < extract_shape(input)[dim]:
-                    slice_1 = graph.call_function(
+                    slice_1 = create_node(
+                        graph,
                         torch.ops.aten.slice_copy.Tensor,
                         args=(
                             input,
@@ -148,13 +152,13 @@ class DecomposeSliceScatter(PassBase):
                             extract_shape(input)[dim],
                             1,
                         ),
+                        origin=node,
                     )
                     slices.append(slice_1)
 
-                concat = graph.call_function(
-                    torch.ops.aten.cat.default, args=(slices, dim)
+                concat = create_node(
+                    graph, torch.ops.aten.cat.default, args=(slices, dim)
                 )
-                # Not set meta for propagating replacing node's meta.
                 node.replace_all_uses_with(concat, propagate_meta=True)
 
             modified = True

tico/passes/fuse_leading_unsqueeze_reshape.py

@@ -20,6 +20,7 @@ from torch.export import ExportedProgram
 from tico.passes import ops
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -84,15 +85,19 @@ class FuseLeadingUnsqueezeReshape(PassBase):
             k = len(back_shape) - len(permute_shape)
             with graph.inserting_before(permute):
                 new_shape = [1] * k + list(reshape_front_size)
-                r_new = graph.call_function(
+                r_new = create_node(
+                    graph,
                     torch.ops.aten.reshape.default,
                     args=(reshape_front_input, new_shape),
+                    origin=reshape_back,
                 )
                 new_p_dims = list(range(k)) + [
                     d + k for d in permute_dims
                 ]  # shift by k
-                p_new = graph.call_function(
-                    torch.ops.aten.permute.default, args=(r_new, new_p_dims)
+                p_new = create_node(
+                    graph,
+                    torch.ops.aten.permute.default,
+                    args=(r_new, new_p_dims),
                 )
 
             reshape_back.replace_all_uses_with(p_new, propagate_meta=True)

tico/passes/legalize_predefined_layout_operators.py

@@ -23,6 +23,7 @@ from torch.export import ExportedProgram
 from tico.serialize.circle_graph import extract_shape
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -124,9 +125,11 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # TODO Introduce a method that inserts permute op.
         # input permute
         with graph.inserting_after(input):
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input, NCHW_to_NHWC),
+                origin=input,
             )
             node.update_arg(node.args.index(input), input_permute)
 
@@ -142,9 +145,11 @@ class LegalizePreDefinedLayoutOperators(PassBase):
             else:
                 assert groups == 1 or groups == input_shape[1]  # Cannot reach here
 
-            weight_permute = graph_module.graph.call_function(
+            weight_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(weight, perm),
+                origin=weight,
             )
             if args.weight.target in [
                 torch.ops.quantized_decomposed.dequantize_per_channel.default,
@@ -171,18 +176,16 @@ class LegalizePreDefinedLayoutOperators(PassBase):
             else:
                 assert groups == 1 or groups == input_shape[1]  # Cannot reach here
 
-            circle_op = graph_module.graph.call_function(
-                legalized_op,
-                args=node.args,
-                kwargs=node.kwargs,
+            circle_op = create_node(
+                graph, legalized_op, args=node.args, kwargs=node.kwargs, origin=node
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            conv_out_permute = graph_module.graph.call_function(
+            conv_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_op, NHWC_to_NCHW),
             )
-            # Not set meta for propagating replacing node's meta.
         node.replace_all_uses_with(conv_out_permute, propagate_meta=True)
 
         logger.debug(f"{node.name} is replaced with {circle_op.name}")
@@ -224,25 +227,29 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         with graph.inserting_after(input):
             # input permute
             NCHW_to_NHWC = [0, 2, 3, 1]
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input, NCHW_to_NHWC),
+                origin=input,
             )
             node.update_arg(node.args.index(input), input_permute)
         with graph.inserting_before(node):
             # circle instnorm
-            circle_instnorm = graph_module.graph.call_function(
+            circle_instnorm = create_node(
+                graph,
                 torch.ops.circle_custom.instance_norm,
                 args=node.args,
                 kwargs=node.kwargs,
+                origin=node,
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            instnorm_out_permute = graph_module.graph.call_function(
+            instnorm_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_instnorm, NHWC_to_NCHW),
             )
-            # Not set meta for propagating replacing node's meta.
         node.replace_all_uses_with(instnorm_out_permute, propagate_meta=True)
 
         logger.debug(f"{node.name} is replaced with {circle_instnorm.name}")
@@ -275,25 +282,25 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # TODO Introduce a method that inserts permute op.
         # input permute
         with graph.inserting_after(input_):
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input_, NCHW_to_NHWC),
+                origin=input_,
             )
             node.update_arg(node.args.index(input_), input_permute)
         with graph.inserting_before(node):
             legalized_op = torch.ops.circle_custom.maxpool2d
-            circle_maxpool2d = graph_module.graph.call_function(
-                legalized_op,
-                args=node.args,
-                kwargs=node.kwargs,
+            circle_maxpool2d = create_node(
+                graph, legalized_op, args=node.args, kwargs=node.kwargs, origin=node
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            maxpool_out_permute = graph_module.graph.call_function(
+            maxpool_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_maxpool2d, NHWC_to_NCHW),
             )
-            # Not set meta for propagating replacing get_item's meta.
         get_item, *_ = node.users.keys()
         get_item.replace_all_uses_with(maxpool_out_permute, propagate_meta=True)
 
@@ -327,21 +334,22 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # TODO Introduce a method that inserts permute op.
         # input permute
         with graph.inserting_after(input_):
-            input_permute = graph_module.graph.call_function(
+            input_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(input_, NCHW_to_NHWC),
+                origin=input_,
             )
             node.update_arg(node.args.index(input_), input_permute)
         with graph.inserting_before(node):
             legalized_op = torch.ops.circle_custom.avgpool2d
-            circle_avgpool2d = graph_module.graph.call_function(
-                legalized_op,
-                args=node.args,
-                kwargs=node.kwargs,
+            circle_avgpool2d = create_node(
+                graph, legalized_op, args=node.args, kwargs=node.kwargs, origin=node
             )
             # output permute
             NHWC_to_NCHW = [0, 3, 1, 2]
-            avgpool_out_permute = graph_module.graph.call_function(
+            avgpool_out_permute = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(circle_avgpool2d, NHWC_to_NCHW),
             )

tico/passes/lower_pow2_to_mul.py

@@ -20,6 +20,7 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -55,7 +56,8 @@ class LowerPow2ToMul(PassBase):
 
             lhs = rhs = in_
             with graph.inserting_after(node):
-                new_mul = graph.call_function(
+                new_mul = create_node(
+                    graph,
                     torch.ops.aten.mul.Tensor,
                     args=(lhs, rhs),
                     kwargs={},

tico/passes/lower_to_resize_nearest_neighbor.py

@@ -22,6 +22,7 @@ from torch.export import ExportedProgram
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -108,18 +109,23 @@ class LowerToResizeNearestNeighbor(PassBase):
         assert expected_shape == list(extract_shape(node))
 
         with graph.inserting_before(node):
-            nchw_to_nhwc = graph.call_function(
-                torch.ops.aten.permute.default, args=(input_tensor, [0, 2, 3, 1])
+            nchw_to_nhwc = create_node(
+                graph,
+                torch.ops.aten.permute.default,
+                args=(input_tensor, [0, 2, 3, 1]),
+                origin=input_tensor,
             )
-            resize_nearest_neighbor = graph.call_function(
+            resize_nearest_neighbor = create_node(
+                graph,
                 torch.ops.circle_custom.resize_nearest_neighbor,
                 args=(nchw_to_nhwc, [len(expected_H_index), len(expected_W_index)]),
+                origin=node,
             )
-            nhwc_to_nchw = graph.call_function(
+            nhwc_to_nchw = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(resize_nearest_neighbor, [0, 3, 1, 2]),
             )
-            # Not set meta for propagating replacing node's meta.
             node.replace_all_uses_with(nhwc_to_nchw, propagate_meta=True)
 
         return resize_nearest_neighbor
@@ -171,18 +177,23 @@ class LowerToResizeNearestNeighbor(PassBase):
             )
 
         with graph.inserting_before(node):
-            nchw_to_nhwc = graph.call_function(
-                torch.ops.aten.permute.default, args=(input_tensor, [0, 2, 3, 1])
+            nchw_to_nhwc = create_node(
+                graph,
+                torch.ops.aten.permute.default,
+                args=(input_tensor, [0, 2, 3, 1]),
+                origin=input_tensor,
             )
-            resize_nearest_neighbor = graph.call_function(
+            resize_nearest_neighbor = create_node(
+                graph,
                 torch.ops.circle_custom.resize_nearest_neighbor,
                 args=(nchw_to_nhwc, [expected_H, expected_W]),
+                origin=node,
             )
-            nhwc_to_nchw = graph.call_function(
+            nhwc_to_nchw = create_node(
+                graph,
                 torch.ops.aten.permute.default,
                 args=(resize_nearest_neighbor, [0, 3, 1, 2]),
             )
-            # Not set meta for propagating replacing node's meta.
             node.replace_all_uses_with(nhwc_to_nchw, propagate_meta=True)
             return resize_nearest_neighbor
 

tico/passes/lower_to_slice.py

@@ -30,7 +30,7 @@ from torch.export import ExportedProgram
 from tico.passes import ops
 from tico.serialize.circle_graph import extract_shape
 from tico.utils import logging
-from tico.utils.graph import is_single_value_tensor
+from tico.utils.graph import create_node, is_single_value_tensor
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_const_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -103,17 +103,22 @@ class LowerSelectCopyToSlice(PassBase):
 
             with graph.inserting_after(node):
                 # slice
-                slice_node = graph.call_function(
-                    torch.ops.aten.slice.Tensor, args=slice_copy_args
+                slice_node = create_node(
+                    graph,
+                    torch.ops.aten.slice.Tensor,
+                    args=slice_copy_args,
+                    origin=node,
                 )
                 node_shape = extract_shape(node)
             with graph.inserting_after(slice_node):
                 # reshape
                 reshape_args = (slice_node, list(node_shape))
-                reshape_node = graph.call_function(
-                    torch.ops.aten.reshape.default, args=reshape_args
+                reshape_node = create_node(
+                    graph,
+                    torch.ops.aten.reshape.default,
+                    args=reshape_args,
                 )
-                node.replace_all_uses_with(reshape_node, propagate_meta=False)
+                node.replace_all_uses_with(reshape_node, propagate_meta=True)
 
             modified = True
             logger.debug(
@@ -196,17 +201,22 @@ class LowerIndexSelectToSlice(PassBase):
 
             with graph.inserting_after(node):
                 # slice
-                slice_node = graph.call_function(
-                    torch.ops.aten.slice.Tensor, args=slice_copy_args
+                slice_node = create_node(
+                    graph,
+                    torch.ops.aten.slice.Tensor,
+                    args=slice_copy_args,
+                    origin=node,
                 )
                 node_shape = extract_shape(node)
             with graph.inserting_after(slice_node):
                 # reshape
                 reshape_args = (slice_node, list(node_shape))
-                reshape_node = graph.call_function(
-                    torch.ops.aten.reshape.default, args=reshape_args
+                reshape_node = create_node(
+                    graph,
+                    torch.ops.aten.reshape.default,
+                    args=reshape_args,
                 )
-                node.replace_all_uses_with(reshape_node, propagate_meta=False)
+                node.replace_all_uses_with(reshape_node, propagate_meta=True)
 
             modified = True
             logger.debug(

tico/passes/remove_redundant_permute.py

@@ -14,7 +14,6 @@
 
 from typing import TYPE_CHECKING
 
-
 if TYPE_CHECKING:
     import torch.fx
 import torch
@@ -23,6 +22,7 @@ from torch.export import ExportedProgram
 from tico.passes import ops
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import is_target_node
@@ -106,8 +106,10 @@ class RemoveRedundantPermutePattern1(PassBase):
             else:
                 with graph.inserting_after(permute2):
                     new_args = (permute1_input, fused_dims)
-                    fused_permute = graph.call_function(
-                        torch.ops.aten.permute.default, args=new_args
+                    fused_permute = create_node(
+                        graph,
+                        torch.ops.aten.permute.default,
+                        args=new_args,
                     )
                     permute2.replace_all_uses_with(fused_permute, propagate_meta=True)
                     logger.debug(f"{permute1.name} and {permute2.name} are fused.")

tico/passes/remove_redundant_reshape.py

@@ -22,6 +22,7 @@ from torch.export import ExportedProgram
 from tico.passes import ops
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 from tico.utils.utils import broadcastable, is_target_node, set_new_meta_val
@@ -369,8 +370,10 @@ class RemoveRedundantReshapePattern4(PassBase):
                 assert isinstance(s, int), type(s)
 
             with graph.inserting_before(reshape1):
-                fused_reshape = graph.call_function(
-                    reshape1.target, (reshape1_input, reshape2_size)
+                fused_reshape = create_node(
+                    graph,
+                    reshape1.target,
+                    (reshape1_input, reshape2_size),
                 )
 
             reshape2.replace_all_uses_with(fused_reshape, propagate_meta=True)

tico/passes/remove_redundant_to_copy.py

@@ -70,6 +70,10 @@ class RemoveRedundantToCopy(PassBase):
                 if input_dtype != target_dtype:
                     continue
 
+            if hasattr(args, "memory_format") and args.memory_format is not None:
+                if args.memory_format != torch.contiguous_format:
+                    continue
+
             node.replace_all_uses_with(input_, propagate_meta=False)
 
             modified = True

tico/passes/restore_linear.py

@@ -16,6 +16,7 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.utils import logging
+from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
 
@@ -74,11 +75,12 @@ class RestoreLinear(PassBase):
 
                 addmm_args = (input, weight, bias)
                 with graph.inserting_after(node):
-                    linear_node = graph.call_function(
-                        torch.ops.aten.linear.default, args=addmm_args
+                    linear_node = create_node(
+                        graph,
+                        torch.ops.aten.linear.default,
+                        args=addmm_args,
                     )
                     node.replace_all_uses_with(linear_node, propagate_meta=True)
-                    graph.erase_node(node)
 
             elif node.target == torch.ops.aten.mm.default:
                 assert len(node.args) == 2
@@ -97,8 +99,8 @@ class RestoreLinear(PassBase):
 
                 mm_args = (input, weight)
                 with graph.inserting_after(node):
-                    linear_node = graph.call_function(
-                        torch.ops.aten.linear.default, args=mm_args
+                    linear_node = create_node(
+                        graph, torch.ops.aten.linear.default, args=mm_args
                     )
                     node.replace_all_uses_with(linear_node, propagate_meta=True)