tico 0.1.0.dev251106__py3-none-any.whl → 0.2.0.dev260122__py3-none-any.whl

tico/__init__.py

@@ -28,8 +28,8 @@ __all__ = [
     "convert_from_pt2",
 ]
 
-# THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev251106"
+# THIS LINE IS AUTOMATICALLY GENERATED
+__version__ = "0.2.0"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/_version.py

@@ -0,0 +1 @@
+__version__ = "0.2.0.dev260122"

tico/passes/convert_conv3d_to_conv2d.py

@@ -0,0 +1,435 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+
+import torch
+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
+from tico.utils.validate_args_kwargs import Conv3DArgs
+
+
+@trace_graph_diff_on_pass
+class ConvertConv3dToConv2d(PassBase):
+    """
+    This pass converts `torch.ops.aten.conv3d` to multiple `torch.ops.aten.conv2d` operations
+
+    [before]                       input(dim=5)         weight(dim=5)
+                                      │                   │
+                                      │                   │
+                                   conv3d<----------------+
+                                      │
+                                      │
+                                   output(dim=5)
+
+    [after]                        input(dim=5)                              weight(dim=5)
+                                      │                                         │
+                                      │                                 ┌───────┴───────┐
+                                      │                                 │ weight slice  │
+                                      │                                 │ (kT times)    │
+                                      │                                 └───────┬───────┘
+                                      │                                         │
+                                      │                                 ┌───────┴───────┐
+                                      │                                 │  squeeze dims │
+                                      │                                 │ (remove dim=2)│
+                                      │                                 └───────┬───────┘
+                                      │                                         │
+                                      │                                 ┌───────┴────────────┐
+                                      │                                 │ weight_2d[0..kT-1] │
+                                      │                                 │ [C_out,C_in,kH,kW] │
+                                      │                                 └───────┬────────────┘
+                                      │                                         │
+                    ┌─────────────────┴──────────────────────────────┐          |
+                    │  temporal padding (if needed)                  │          |
+                    │  ┌────────────┐  ┌────────────┐  ┌───────────┐ │          |
+                    │  │ zeros      │  │  input     │  │zeros      │ │          |
+                    │  │ [N,C,p,H,W]│  │ [N,C,T,H,W]│  │[N,C,p,H,W]│ │          |
+                    │  └────┬───────┘  └────┬───────┘  └────┬──────┘ │          |
+                    │       └───────────┼───┴───────────────┘        │          |
+                    │                   │                            │          |
+                    │           ┌───────┴───────┐                    │          |
+                    │           │     cat       │                    │          |
+                    │           │ (dim=2)       │                    │          |
+                    │           └───────┬───────┘                    │          |
+                    │                   │                            │          |
+                    │           ┌───────┴───────┐                    │          |
+                    │           │  padded_input │                    │          |
+                    │           │ [N,C,T+2p,H,W]│                    │          |
+                    │           └───────┬───────┘                    │          |
+                    └───────────────────┼────────────────────────────┘          |
+                                        │                                       |
+                    ┌───────────────────┴───────────────────────────────┐       |
+                    │           Temporal Processing Loop                │       |
+                    │  ┌────────────────────────────────────────────┐   │       |
+                    │  │ For t_out = 0..T_out-1:                    │   │       |
+                    │  │   For i = 0..kT-1:                         │   │       |
+                    │  │     t_idx = t_out*stride[0] + i*dilation[0]│   │       |
+                    │  │     ┌─────────────────────────┐            │   │       |
+                    │  │     │ slice input[t_idx]      │            │   │       |
+                    │  │     │ [N,C,H,W]               │            │   │       |
+                    │  │     └─────────┬───────────────┘            │   │       |
+                    │  │               │                            │   │       |
+                    │  │     ┌─────────┴───────────────┐            │   │       |
+                    │  │     │ squeeze dims            │            │   │       |
+                    │  │     │ [N,C,H,W]               │            │   │       |
+                    │  │     └─────────┬───────────────┘            │   │       |
+                    │  │               │                            │   │       |
+                    │  │     ┌─────────┴───────────────┐            │   │       |
+                    │  │     │ conv2d(input,weight)    │            │   │───────┘
+                    │  │     │ [N,C_out,H_out,W_out]   │            │   │
+                    │  │     └─────────┬───────────────┘            │   │
+                    │  │               │                            │   │
+                    │  │     ┌─────────┴───────────────┐            │   │
+                    │  │     │ where(valid_mask,       │            │   │
+                    │  │     │      conv2d, zeros)     │            │   │
+                    │  │     └─────────┬───────────────┘            │   │
+                    │  │               │                            │   │
+                    │  │     ┌─────────┴───────────────┐            │   │
+                    │  │     │ accumulate (add)        │            │   │
+                    │  │     └─────────┬───────────────┘            │   │
+                    │  └───────────────┼────────────────────────────┘   │
+                    │                  │                                │
+                    │           ┌──────┴───────────┐                    │
+                    │           │ add bias (if any)│                    │
+                    │           └───────┬──────────┘                    │
+                    │                   │                               │
+                    │           ┌───────┴──────────┐                    │
+                    │           │ unsqueeze (dim=2)│                    │
+                    │           └───────┬──────────┘                    │
+                    └───────────────────┼───────────────────────────────┘
+                                        │
+                    ┌───────────────────┴───────────────────────┐
+                    │           cat (dim=2)                     │
+                    │           [N,C_out,T_out,H_out,W_out]     │
+                    └───────────────────┬───────────────────────┘
+                                        │
+                                   output(dim=5)
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def _parse_3d_padding(self, padding, kernel_size):
+        """
+        Parse 3D padding parameter and return (temporal, H, W) tuple.
+
+        Args:
+            padding: Can be str ('same', 'valid'), int, list, or tuple
+            kernel_size: 3D kernel size (kT, kH, kW)
+
+        Returns:
+            Tuple of 3 padding values: (temporal_padding, H_padding, W_padding)
+        """
+        if isinstance(padding, str):
+            if padding == "same":
+                # For 'same' padding, use kernel_size // 2
+                if isinstance(kernel_size, int):
+                    return kernel_size // 2, kernel_size // 2, kernel_size // 2
+                else:
+                    return kernel_size[0] // 2, kernel_size[1] // 2, kernel_size[2] // 2
+            elif padding == "valid":
+                return 0, 0, 0
+            else:
+                raise NotYetSupportedError(f"Unsupported padding string: {padding}")
+        elif isinstance(padding, (list, tuple)):
+            if len(padding) == 1:
+                return padding[0], padding[0], padding[0]
+            elif len(padding) == 3:
+                return padding[0], padding[1], padding[2]
+            else:
+                raise NotYetSupportedError(f"Unsupported padding format: {padding}")
+        else:  # int
+            return padding, padding, padding
+
+    def convert(self, exported_program: ExportedProgram, node: torch.fx.Node) -> bool:
+        logger = logging.getLogger(__name__)
+        modified = False
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+
+        # Extract conv3d arguments
+        args = Conv3DArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+
+        input = args.input
+        weight = args.weight
+        bias = args.bias
+        stride = args.stride
+        padding = args.padding
+        dilation = args.dilation
+        groups = args.groups
+
+        input_shape = extract_shape(input)
+        weight_shape = extract_shape(weight)
+
+        if not (len(input_shape) == 5):
+            raise NotYetSupportedError(
+                f"Only support 5D input tensor: node's input shape: {input_shape}"
+            )
+
+        if not (len(weight_shape) == 5):
+            raise NotYetSupportedError(
+                f"Only support 5D weight tensor: node's weight shape: {weight_shape}"
+            )
+
+        N, C_in, T_in, H_in, W_in = input_shape
+        C_out, C_in_weight, kT, kH, kW = weight_shape
+
+        temporal_padding, h_padding, w_padding = self._parse_3d_padding(
+            padding, (kT, kH, kW)
+        )
+
+        # Calculate output dimensions
+        T_out = (T_in + 2 * temporal_padding - dilation[0] * (kT - 1) - 1) // stride[
+            0
+        ] + 1
+
+        H_out = (H_in + 2 * h_padding - dilation[1] * (kH - 1) - 1) // stride[1] + 1
+        W_out = (W_in + 2 * w_padding - dilation[2] * (kW - 1) - 1) // stride[2] + 1
+
+        # Find the next node after conv3d
+        next_node = node.next
+        if next_node is None:
+            # If no next node, find the output node
+            for n in graph.nodes:
+                if n.op == "output":
+                    next_node = n
+                    break
+
+        if next_node is None:
+            raise RuntimeError("Could not find insertion point for temporal outputs")
+
+        # Create all nodes before the next node in one go
+        with graph.inserting_before(next_node):
+            # Step 1: Create weight_2d layers first (they depend only on weight)
+            weight_2d_layers = []
+            for t in range(kT):
+                # Slice weight for temporal dimension t: [C_out, C_in, t, kH, kW] -> [C_out, C_in, kH, kW]
+                weight_slice = create_node(
+                    graph,
+                    torch.ops.aten.slice.Tensor,
+                    args=(weight, 2, t, t + 1, 1),
+                    origin=weight,
+                )
+
+                # Remove temporal dimension: [C_out, C_in, 1, kH, kW] -> [C_out, C_in, kH, kW]
+                weight_2d = create_node(
+                    graph,
+                    torch.ops.aten.squeeze.dims,
+                    args=(weight_slice, [2]),
+                    origin=weight_slice,
+                )
+                weight_2d_layers.append(weight_2d)
+
+            # Step 2: Create padded input (if needed) using cat
+            if temporal_padding > 0:
+                # Create zero padding: [N, C, padding, H, W]
+                zero_padding = create_node(
+                    graph,
+                    torch.ops.aten.zeros.default,
+                    args=([N, C_in, temporal_padding, H_in, W_in],),
+                    kwargs={
+                        "dtype": input.meta.get("dtype", torch.float32),
+                        "device": input.meta.get("device", "cpu"),
+                    },
+                    origin=input,
+                )
+
+                # Cat: [zero_padding, input, zero_padding] -> [N, C, T+2*padding, H, W]
+                padded_input = create_node(
+                    graph,
+                    torch.ops.aten.cat.default,
+                    args=([zero_padding, input, zero_padding], 2),
+                    origin=input,
+                )
+                T_padded = T_in + 2 * temporal_padding
+            else:
+                padded_input = input
+                T_padded = T_in
+
+            # Step 3: Process each temporal output position
+            temporal_outputs = []
+            for t_out in range(T_out):
+                # Calculate input time position
+                t_in = t_out * stride[0]
+
+                # Initialize accumulator for this temporal position
+                acc = None
+
+                for i, weight_2d in enumerate(weight_2d_layers):
+                    # Calculate actual time index with dilation
+                    t_idx = t_in + i * dilation[0]
+
+                    # Create constant for time index
+                    t_idx_const = create_node(
+                        graph,
+                        torch.ops.aten.scalar_tensor.default,
+                        args=(t_idx,),
+                        kwargs={"dtype": torch.int64},
+                        origin=node,
+                    )
+
+                    # Create constant for T_padded
+                    t_padded_const = create_node(
+                        graph,
+                        torch.ops.aten.scalar_tensor.default,
+                        args=(T_padded,),
+                        kwargs={"dtype": torch.int64},
+                        origin=node,
+                    )
+
+                    # Check if t_idx < T_padded
+                    valid_mask = create_node(
+                        graph,
+                        torch.ops.aten.lt.Tensor,
+                        args=(t_idx_const, t_padded_const),
+                        origin=node,
+                    )
+
+                    # Slice input at time t_idx: [N, C_in, T_padded, H_in, W_in] -> [N, C_in, H_in, W_in]
+                    input_slice = create_node(
+                        graph,
+                        torch.ops.aten.slice.Tensor,
+                        args=(padded_input, 2, t_idx, t_idx + 1, 1),
+                        origin=padded_input,
+                    )
+
+                    # Remove temporal dimension: [N, C_in, 1, H_in, W_in] -> [N, C_in, H_in, W_in]
+                    input_2d = create_node(
+                        graph,
+                        torch.ops.aten.squeeze.dims,
+                        args=(input_slice, [2]),
+                        origin=input_slice,
+                    )
+
+                    # Create conv2d operation with proper input
+                    conv2d = create_node(
+                        graph,
+                        torch.ops.aten.conv2d.default,
+                        args=(
+                            input_2d,  # input is now available
+                            weight_2d,
+                            None,  # bias = False
+                            [stride[1], stride[2]],
+                            [h_padding, w_padding],
+                            [dilation[1], dilation[2]],
+                            groups,
+                        ),
+                        origin=node,
+                    )
+
+                    # Create zero tensor with calculated shape
+                    # conv2d output shape: [N, C_out, H_out, W_out]
+                    zero_tensor = create_node(
+                        graph,
+                        torch.ops.aten.zeros.default,
+                        args=([N, C_out, H_out, W_out],),
+                        kwargs={
+                            "dtype": input.meta.get("dtype", torch.float32),
+                            "device": input.meta.get("device", "cpu"),
+                        },
+                        origin=conv2d,
+                    )
+
+                    # Apply conditional execution
+                    conv2d_masked = create_node(
+                        graph,
+                        torch.ops.aten.where.self,
+                        args=(valid_mask, conv2d, zero_tensor),
+                        origin=conv2d,
+                    )
+
+                    if acc is None:
+                        # First temporal slice
+                        acc = conv2d_masked
+                    else:
+                        # Add subsequent temporal slices
+                        acc = create_node(
+                            graph,
+                            torch.ops.aten.add.Tensor,
+                            args=(acc, conv2d_masked),
+                            origin=acc,
+                        )
+
+                # Add bias if present
+                if bias is not None:
+                    bias_reshaped = create_node(
+                        graph,
+                        torch.ops.aten.reshape.default,
+                        args=(bias, [1, C_out, 1, 1]),
+                        origin=bias,
+                    )
+                    acc = create_node(
+                        graph,
+                        torch.ops.aten.add.Tensor,
+                        args=(acc, bias_reshaped),
+                        origin=acc,
+                    )
+
+                temporal_outputs.append(acc)
+
+            # Step 4: Stack temporal outputs using cat instead of stack
+            # First, unsqueeze each temporal output to add the time dimension
+            unsqueezed_outputs = []
+            for i, temp_output in enumerate(temporal_outputs):
+                # Add time dimension: [N, C_out, H_out, W_out] -> [N, C_out, 1, H_out, W_out]
+                unsqueezed = create_node(
+                    graph,
+                    torch.ops.aten.unsqueeze.default,
+                    args=(temp_output, 2),
+                    origin=temp_output,
+                )
+                unsqueezed_outputs.append(unsqueezed)
+
+            # Cat along time dimension: [N, C_out, T_out, H_out, W_out]
+            stacked_output = create_node(
+                graph,
+                torch.ops.aten.cat.default,
+                args=(unsqueezed_outputs, 2),
+                origin=node,
+            )
+
+        # Replace the original node
+        node.replace_all_uses_with(stacked_output, propagate_meta=False)
+        logger.debug(f"{node.name} is replaced with conv2d decomposition")
+        modified = True
+
+        return modified
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        target_conv_op = [torch.ops.aten.conv3d.default, torch.ops.aten.conv3d.padding]
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+
+        modified = False
+
+        # Process all Conv3D nodes in forward pass order
+        for node in graph.nodes:
+            if not is_target_node(node, target_conv_op):
+                continue
+            modified |= self.convert(exported_program, node)
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/passes/convert_sym_size_to_circle_shape.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+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
+
+
+@trace_graph_diff_on_pass
+class ConvertSymSizeToCircleShape(PassBase):
+    """
+    This pass converts torch.ops.aten.sym_size.int operations to circle_custom::shape.
+
+    The circle_custom::shape operator allows preserving dynamic shape information
+    in the Circle model. This is essential for models with dynamic batch sizes or other dynamic dimensions.
+
+    Example:
+        Before: %sym_size_int_1 = call_function[target=torch.ops.aten.sym_size.int](args=(%x, 0))
+        After:  %shape_0 = call_function[target=torch.ops.circle_custom.shape](args=(%x,))
+                %slice_0 = call_function[target=torch.ops.aten.slice.Tensor](args=(%shape_0, 0, 0, 1, 1))
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+
+            if node.target == torch.ops.aten.sym_size.int:
+                # sym_size.int has args: (input, dim)
+                input_tensor = node.args[0]
+                dim = node.args[1]
+
+                # Create circle_custom::shape node
+                with graph.inserting_after(node):
+                    shape_node = create_node(
+                        graph,
+                        torch.ops.circle_custom.shape,
+                        args=(input_tensor,),
+                    )
+
+                # Set metadata for shape_node
+                if "val" in input_tensor.meta:
+                    input_val = input_tensor.meta["val"]
+                    rank = len(input_val.shape)
+                    # shape output is a 1D tensor of size rank, dtype int32
+                    # We use a real tensor here as a placeholder for metadata
+                    shape_node.meta["val"] = torch.zeros(rank, dtype=torch.int32)
+
+                # Extract the specific dimension using slice
+                with graph.inserting_after(shape_node):
+                    slice_node = create_node(
+                        graph,
+                        torch.ops.aten.slice.Tensor,
+                        args=(shape_node, 0, dim, dim + 1, 1),
+                    )
+                    # slice output is 1D tensor of size 1
+                    slice_node.meta["val"] = torch.zeros(1, dtype=torch.int32)
+
+                # Replace all uses
+                node.replace_all_uses_with(slice_node, propagate_meta=False)
+                modified = True
+
+                logger.debug(
+                    f"Converted {node.name} (sym_size.int) to {shape_node.name} (circle_custom::shape) + {slice_node.name} (slice)"
+                )
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/passes/decompose_batch_norm.py

@@ -115,7 +115,7 @@ class DecomposeBatchNorm(PassBase):
                 continue
 
             input_shape = extract_shape(input_)
-            assert len(input_shape) == 4
+            assert len(input_shape) >= 2, len(input_shape)
             C = input_shape[1]
 
             weight_value = (
@@ -145,11 +145,15 @@ class DecomposeBatchNorm(PassBase):
             # Calculate constants for mul and add
             mul_const = weight_value / torch.sqrt(var_value + eps)
             add_const = bias_value - (mul_const * mean_value)
-            # N, C, H, W
+
+            # Make sure channel count matches
             assert len(mul_const) == len(add_const) == C
-            # reshape along with channel dimension
-            mul_const = mul_const.view(1, mul_const.shape[0], 1, 1)
-            add_const = add_const.view(1, add_const.shape[0], 1, 1)
+
+            # Build a broadcastable shape like (1, C, 1, ...)
+            view_shape = [1] * len(input_shape)
+            view_shape[1] = C
+            mul_const = mul_const.view(*view_shape)
+            add_const = add_const.view(*view_shape)
 
             # Placeholder nodes must be the first N nodes in the nodes list of a graph.
             # Therefore, insert the newly created placeholders at the start of the node list.

tico/passes/lower_copy.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+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.validate_args_kwargs import CopyArgs
+
+
+@trace_graph_diff_on_pass
+class LowerCopy(PassBase):
+    """
+    This pass lowers `aten.copy.default` to simpler broadcast operations.
+
+    - If src and dst shapes are the same, the copy is redundant and folded away.
+    - If src and dst shapes differ, it's replaced with expand (broadcast).
+
+    This simplifies serialization by handling copy logic at the pass level.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+
+        for node in graph.nodes:
+            if not node.op == "call_function":
+                continue
+
+            if node.target != torch.ops.aten.copy.default:
+                continue
+
+            args = CopyArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            dst = args.dst
+            src = args.src
+
+            dst_shape = list(extract_shape(dst))
+            src_shape = list(extract_shape(src))
+
+            # Case 1: Same shape - copy is redundant, just use src
+            if dst_shape == src_shape:
+                logger.debug(
+                    f"{node.name}: Same shape {dst_shape}, replacing with src directly"
+                )
+                node.replace_all_uses_with(src, propagate_meta=False)
+                modified = True
+                continue
+
+            # Case 2: Different shapes - need expand
+            logger.debug(
+                f"{node.name}: Different shapes src={src_shape} dst={dst_shape}, "
+                f"inserting expand"
+            )
+
+            with graph.inserting_before(node):
+                expand_node = create_node(
+                    graph,
+                    torch.ops.aten.expand.default,
+                    args=(src, dst_shape),
+                )
+
+            node.replace_all_uses_with(expand_node, propagate_meta=True)
+            modified = True
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)