tico 0.1.0.dev250610__py3-none-any.whl → 0.1.0.dev250611__py3-none-any.whl

tico/passes/decompose_slice_scatter.py

@@ -25,7 +25,7 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 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
+from tico.utils.utils import enforce_type, is_target_node
 
 
 @trace_graph_diff_on_pass
@@ -82,9 +82,7 @@ class DecomposeSliceScatter(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if node.op != "call_function":
-                continue
-            if node.target != torch.ops.aten.slice_scatter.default:
+            if not is_target_node(node, torch.ops.aten.slice_scatter.default):
                 continue
 
             @enforce_type

tico/passes/extract_dtype_kwargs.py

@@ -23,6 +23,7 @@ from torch.utils import _pytree as pytree
 from tico.utils import logging
 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
 
 
 def _extract_to_output(node: torch.fx.Node, graph: torch.fx.Graph) -> bool:
@@ -107,7 +108,7 @@ class ExtractDtypeKwargsPass(PassBase):
         graph: torch.fx.Graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function" or node.target not in self.target_ops:
+            if not is_target_node(node, list(self.target_ops.keys())):
                 continue
             if "dtype" not in node.kwargs:
                 continue

tico/passes/fuse_leading_unsqueeze_reshape.py

@@ -22,7 +22,7 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
-from tico.utils.utils import is_single_use_target_node
+from tico.utils.utils import is_target_node
 from tico.utils.validate_args_kwargs import PermuteArgs, ReshapeArgs
 
 
@@ -55,20 +55,17 @@ class FuseLeadingUnsqueezeReshape(PassBase):
         graph = gm.graph
         modified = False
         for reshape_back in graph.nodes:
-            if (
-                reshape_back.op != "call_function"
-                or reshape_back.target not in ops.aten.reshape
-            ):
+            if not is_target_node(reshape_back, ops.aten.reshape):
                 continue
             reshape_back_args = ReshapeArgs(*reshape_back.args, **reshape_back.kwargs)  # type: ignore[arg-type]
             permute, reshape_back_size = reshape_back_args.input, reshape_back_args.size
 
-            if not is_single_use_target_node(permute, ops.aten.permute):
+            if not is_target_node(permute, ops.aten.permute):
                 continue
             permute_args = PermuteArgs(*permute.args, **permute.kwargs)  # type: ignore[arg-type]
             reshape_front, permute_dims = permute_args.input, permute_args.dims
 
-            if not is_single_use_target_node(reshape_front, ops.aten.reshape):
+            if not is_target_node(reshape_front, ops.aten.reshape):
                 continue
             reshape_front_args = ReshapeArgs(*reshape_front.args, **reshape_front.kwargs)  # type: ignore[arg-type]
             reshape_front_input, reshape_front_size = (

tico/passes/fuse_redundant_reshape_to_mean.py

@@ -25,6 +25,7 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 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
 
 
 @trace_graph_diff_on_pass
@@ -45,10 +46,7 @@ class FuseRedundantReshapeToMean(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if node.target != torch.ops.aten.mean.dim:
+            if not is_target_node(node, torch.ops.aten.mean.dim):
                 continue
 
             # If mean is being used in other nodes, do not fuse it.
@@ -56,7 +54,7 @@ class FuseRedundantReshapeToMean(PassBase):
                 continue
 
             user_node = next(iter(node.users))
-            if user_node.target not in ops.aten.reshape:
+            if not is_target_node(user_node, ops.aten.reshape):
                 continue
 
             mean_args, mean_kwargs = pytree.tree_map_only(

tico/passes/legalize_causal_mask_value.py

@@ -20,10 +20,10 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.passes import ops
-
 from tico.utils import logging
 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 AddTensorArgs
 
 
@@ -53,14 +53,9 @@ class LegalizeCausalMaskValue(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.add:
+            if not is_target_node(node, ops.aten.add):
                 continue
 
-            assert len(node.args) == 2
-
             args = AddTensorArgs(*node.args, **node.kwargs)
             input = args.input
             other = args.other

tico/passes/legalize_predefined_layout_operators.py

@@ -25,6 +25,7 @@ from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 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 (
     AvgPool2dArgs,
     Conv2DArgs,
@@ -363,11 +364,9 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, list(target_to_legalize_func.keys())):
                 continue
 
-            if node.target not in target_to_legalize_func:
-                continue
             modified |= target_to_legalize_func[node.target](exported_program, node)
 
         graph.eliminate_dead_code()

tico/passes/lower_pow2_to_mul.py

@@ -22,6 +22,8 @@ from torch.export import ExportedProgram
 from tico.utils import logging
 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 PowTensorScalarArgs
 
 
 @trace_graph_diff_on_pass
@@ -42,15 +44,11 @@ class LowerPow2ToMul(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, torch.ops.aten.pow.Tensor_Scalar):
                 continue
 
-            if node.target != torch.ops.aten.pow.Tensor_Scalar:
-                continue
-
-            assert len(node.args) == 2, len(node.args)
-            in_, exp = node.args
-            assert isinstance(in_, torch.fx.Node), type(in_)
+            args = PowTensorScalarArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            in_, exp = args.input, args.exponent
 
             if exp != 2:
                 continue

tico/passes/lower_to_resize_nearest_neighbor.py

@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
+from typing import Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch.fx
-from typing import Optional
-
 import torch
 from torch.export import ExportedProgram
 
@@ -26,6 +24,7 @@ from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 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 IndexArgs, UpsampleNearest2DVecArgs
 
 
@@ -194,13 +193,10 @@ class LowerToResizeNearestNeighbor(PassBase):
         graph_module = exported_program.graph_module
         graph = graph_module.graph
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if node.target not in [
-                torch.ops.aten.index.Tensor,
-                torch.ops.aten.upsample_nearest2d.vec,
-            ]:
+            if not is_target_node(
+                node,
+                [torch.ops.aten.index.Tensor, torch.ops.aten.upsample_nearest2d.vec],
+            ):
                 continue
 
             resize_nearest_neighbor = None

tico/passes/lower_to_slice.py

@@ -28,13 +28,12 @@ from torch._export.utils import (
 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.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_const_diff_on_pass
+from tico.utils.utils import is_target_node
 from tico.utils.validate_args_kwargs import IndexSelectArgs, SelectCopyIntArgs
 
 
@@ -85,10 +84,7 @@ class LowerSelectCopyToSlice(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.select:
+            if not is_target_node(node, ops.aten.select):
                 continue
 
             args = SelectCopyIntArgs(*node.args, **node.kwargs)
@@ -163,11 +159,9 @@ class LowerIndexSelectToSlice(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, ops.aten.index_select):
                 continue
 
-            if not node.target in ops.aten.index_select:
-                continue
             args = IndexSelectArgs(*node.args, **node.kwargs)
             input = args.input
             dim = args.dim

tico/passes/merge_consecutive_cat.py

@@ -18,6 +18,7 @@ from tico.passes import ops
 from tico.utils import logging
 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 CatArgs
 
 
@@ -37,10 +38,7 @@ class MergeConsecutiveCat(PassBase):
         graph = graph_module.graph
         modified = False
         for cat in graph.nodes:
-            if not cat.op == "call_function":
-                continue
-
-            if not cat.target in ops.aten.cat:
+            if not is_target_node(cat, ops.aten.cat):
                 continue
 
             args = CatArgs(*cat.args, **cat.kwargs)  # type: ignore[arg-type]

tico/passes/remove_nop.py

@@ -23,6 +23,7 @@ from tico.passes import ops
 from tico.utils import logging
 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
 
 
 @trace_graph_diff_on_pass
@@ -51,11 +52,9 @@ class RemoveNop(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, RemoveNop.target_ops):
                 continue
 
-            if not node.target in RemoveNop.target_ops:
-                continue
             # TODO Consider memory format
             if node.target in ops.aten.clone and "memory_format" in node.kwargs:
                 if node.kwargs["memory_format"] not in [

tico/passes/remove_redundant_assert_nodes.py

@@ -17,6 +17,7 @@ from torch.export import ExportedProgram
 
 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
 
 
 assert_node_targets = [
@@ -39,7 +40,7 @@ class RemoveRedundantAssertionNodes(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if node.op == "call_function" and node.target in assert_node_targets:
+            if is_target_node(node, assert_node_targets):
                 graph.erase_node(node)
                 modified = True
 

tico/passes/remove_redundant_expand.py

@@ -24,6 +24,8 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 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 ExpandArgs
 
 
 @trace_graph_diff_on_pass
@@ -42,17 +44,11 @@ class RemoveRedundantExpand(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, ops.aten.expand):
                 continue
 
-            if not node.target in ops.aten.expand:
-                continue
-
-            assert len(node.args) == 2
-
-            input, size = list(node.args)
-            assert isinstance(input, torch.fx.Node), type(input)
-            assert isinstance(size, list), type(size)
+            args = ExpandArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            input, size = args.input, args.size
 
             input_shape = extract_shape(input)
             if list(input_shape) != size:

tico/passes/remove_redundant_permute.py

@@ -14,16 +14,18 @@
 
 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, extract_stride
+from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 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 PermuteArgs
 
 
@@ -75,16 +77,15 @@ class RemoveRedundantPermutePattern1(PassBase):
         graph = graph_module.graph
         modified = False
         for permute2 in graph.nodes:
-            if not permute2.op == "call_function":
-                continue
-            if not permute2.target in ops.aten.permute:
+            if not is_target_node(permute2, ops.aten.permute):
                 continue
+
             if len(permute2.users) != 1:
                 continue
             permute2_args = PermuteArgs(*permute2.args, **permute2.kwargs)  # type: ignore[arg-type]
             permute1, permute2_dims = permute2_args.input, permute2_args.dims
 
-            if not permute1.target in ops.aten.permute:
+            if not is_target_node(permute1, ops.aten.permute):
                 continue
             if len(permute1.users) != 1:
                 continue

tico/passes/remove_redundant_reshape.py

@@ -24,7 +24,7 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
-from tico.utils.utils import broadcastable, set_new_meta_val
+from tico.utils.utils import broadcastable, is_target_node, set_new_meta_val
 from tico.utils.validate_args_kwargs import (
     AddTensorArgs,
     PermuteArgs,
@@ -67,12 +67,10 @@ class RemoveRedundantReshapePattern1(PassBase):
         graph = graph_module.graph
         modified = False
         for reshape1 in graph.nodes:
-            if not reshape1.op == "call_function":
-                continue
-
             ### first reshape
-            if not reshape1.target in ops.aten.reshape:
+            if not is_target_node(reshape1, ops.aten.reshape):
                 continue
+
             # Assumes that other node do not use ops in the pattern for simplisity.
             if len(reshape1.users) != 1:
                 continue
@@ -86,7 +84,7 @@ class RemoveRedundantReshapePattern1(PassBase):
 
             ### permute
             permute = next(iter(reshape1.users))
-            if not permute.target in ops.aten.permute:
+            if not is_target_node(permute, ops.aten.permute):
                 continue
             if len(permute.users) != 1:
                 continue
@@ -98,14 +96,14 @@ class RemoveRedundantReshapePattern1(PassBase):
 
             ### mul
             mul = next(iter(permute.users))
-            if not mul.target in RemoveRedundantReshapePattern1.mul_ops:
+            if not is_target_node(mul, RemoveRedundantReshapePattern1.mul_ops):
                 continue
             if len(mul.users) != 1:
                 continue
 
             ### second reshape
             reshape2 = next(iter(mul.users))
-            if not reshape2.target in ops.aten.reshape:
+            if not is_target_node(reshape2, ops.aten.reshape):
                 continue
             if len(reshape2.users) != 1:
                 continue
@@ -153,11 +151,8 @@ class RemoveRedundantReshapePattern2(PassBase):
         graph = graph_module.graph
         modified = False
         for reshape1 in graph.nodes:
-            if not reshape1.op == "call_function":
-                continue
-
             ### first reshape
-            if not reshape1.target in ops.aten.reshape:
+            if not is_target_node(reshape1, ops.aten.reshape):
                 continue
             if len(reshape1.users) != 1:
                 continue
@@ -171,7 +166,7 @@ class RemoveRedundantReshapePattern2(PassBase):
 
             ### permute
             permute = next(iter(reshape1.users))
-            if not permute.target in ops.aten.permute:
+            if not is_target_node(permute, ops.aten.permute):
                 continue
             if len(permute.users) != 1:
                 continue
@@ -183,7 +178,7 @@ class RemoveRedundantReshapePattern2(PassBase):
 
             ### second reshape
             reshape2 = next(iter(permute.users))
-            if not reshape2.target in ops.aten.reshape:
+            if not is_target_node(reshape2, ops.aten.reshape):
                 continue
             if len(reshape2.users) != 1:
                 continue
@@ -239,20 +234,14 @@ class RemoveRedundantReshapePattern3(PassBase):
         graph = graph_module.graph
         modified = False
         for reshape_1 in graph.nodes:
-            assert isinstance(reshape_1, torch.fx.Node), type(reshape_1)
             # reshape_1
-            if not reshape_1.op == "call_function":
-                continue
-            if not reshape_1.target in ops.aten.reshape:
+            if not is_target_node(reshape_1, ops.aten.reshape):
                 continue
             reshape_1_args = ReshapeArgs(*reshape_1.args, **reshape_1.kwargs)  # type: ignore[arg-type]
-            softmax, reshape_1_size = reshape_1_args.input, reshape_1_args.size
+            softmax = reshape_1_args.input
 
             # softmax
-            assert isinstance(softmax, torch.fx.Node), type(softmax)
-            if not softmax.op == "call_function":
-                continue
-            if not softmax.target in ops.aten.softmax:
+            if not is_target_node(softmax, ops.aten.softmax):
                 continue
             if softmax.target == torch.ops.aten._softmax.default:
                 softmax_args = SoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type, assignment]
@@ -354,10 +343,9 @@ class RemoveRedundantReshapePattern4(PassBase):
         modified = False
         for reshape1 in graph.nodes:
             # reshape_1
-            if not reshape1.op == "call_function":
-                continue
-            if not reshape1.target in ops.aten.reshape:
+            if not is_target_node(reshape1, ops.aten.reshape):
                 continue
+
             reshape1_args = ReshapeArgs(*reshape1.args, **reshape1.kwargs)  # type: ignore[arg-type]
             reshape1_input, size = reshape1_args.input, reshape1_args.size
             assert isinstance(reshape1_input, torch.fx.Node), type(reshape1_input)
@@ -370,12 +358,10 @@ class RemoveRedundantReshapePattern4(PassBase):
 
             # reshape_2
             reshape2 = next(iter(reshape1.users))
-            if not reshape2.op == "call_function":
+            if not is_target_node(reshape2, ops.aten.reshape):
                 continue
-            if not reshape2.target in ops.aten.reshape:
-                continue
-            reshape2_args = ReshapeArgs(*reshape2.args, **reshape2.kwargs)  # type: ignore[arg-type]
 
+            reshape2_args = ReshapeArgs(*reshape2.args, **reshape2.kwargs)  # type: ignore[arg-type]
             reshape2_input, reshape2_size = reshape2_args.input, reshape2_args.size
             assert isinstance(reshape2_input, torch.fx.Node), type(reshape2_input)
             assert isinstance(reshape2_size, list), type(reshape2_size)
@@ -420,10 +406,7 @@ class RemoveRedundantReshapePattern5(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.reshape:
+            if not is_target_node(node, ops.aten.reshape):
                 continue
 
             args = ReshapeArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]

tico/passes/remove_redundant_slice.py

@@ -19,6 +19,7 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 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 SliceArgs
 
 
@@ -38,10 +39,7 @@ class RemoveRedundantSlice(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.slice:
+            if not is_target_node(node, ops.aten.slice):
                 continue
 
             args = SliceArgs(*node.args, **node.kwargs)

tico/passes/remove_redundant_to_copy.py

@@ -22,6 +22,7 @@ from tico.serialize.circle_mapping import extract_torch_dtype
 from tico.utils import logging
 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 ToCopyArgs, ToDtypeArgs, ToDtypeLayoutArgs
 
 
@@ -41,10 +42,7 @@ class RemoveRedundantToCopy(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.to_copy:
+            if not is_target_node(node, ops.aten.to_copy):
                 continue
 
             args: Union[ToCopyArgs, ToDtypeArgs, ToDtypeLayoutArgs]

tico/passes/segment_index_select.py

@@ -34,6 +34,7 @@ from tico.utils import logging
 from tico.utils.graph import add_placeholder, 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
 from tico.utils.validate_args_kwargs import IndexSelectArgs
 
 
@@ -78,10 +79,7 @@ class SegmentIndexSelectConst(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if not node.target in ops.aten.index_select:
+            if not is_target_node(node, ops.aten.index_select):
                 continue
 
             args = IndexSelectArgs(*node.args, **node.kwargs)

tico/serialize/operators/op_where.py

@@ -52,12 +52,12 @@ class WhereVisitor(NodeVisitor):
 
         result_true_dtype = (
             extract_torch_dtype(input)
-            if isinstance(input, torch.fx.node.Node)
+            if isinstance(input, torch.fx.Node)
             else input.dtype  # type: ignore[union-attr]
         )
         result_false_dtype = (
             extract_torch_dtype(other)
-            if isinstance(other, torch.fx.node.Node)
+            if isinstance(other, torch.fx.Node)
             else other.dtype  # type: ignore[union-attr]
         )
 

tico/utils/utils.py

@@ -331,6 +331,7 @@ def get_quant_dtype(qmin: int, qmax: int):
     """
     known_ranges = {
         (-32768, 32767): "int16",
+        (-32767, 32767): "int16",
         (0, 65535): "uint16",
         (-128, 127): "int8",
         (0, 255): "uint8",
@@ -380,19 +381,18 @@ def broadcastable(
     return True
 
 
-def is_single_use_target_node(
+def is_target_node(
     node: torch.fx.Node, target_ops: list[torch._ops.OpOverload] | torch._ops.OpOverload
 ):
     """
-    Check whether a given node is a `call_function` node that matches one of the specified targets
-    and is used by only one other node.
+    Check whether a given node is a `call_function` node that matches one of the specified targets.
 
     Args:
         node (torch.fx.Node): The node to check.
         target_ops (Iterable[Callable]): A list or set of target operations to match (e.g., ops.aten.reshape).
 
     Returns:
-        bool: True if the node is a call_function, its target is in `target_ops`, and it has exactly one user.
+        bool: True if the node is a call_function, its target is in `target_ops`.
     """
     if not isinstance(target_ops, list):
         target_ops = [target_ops]
@@ -402,7 +402,5 @@ def is_single_use_target_node(
         return False
     if node.target not in target_ops:
         return False
-    if len(node.users) != 1:
-        return False
 
     return True

{tico-0.1.0.dev250610.dist-info → tico-0.1.0.dev250611.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250610
+Version: 0.1.0.dev250611
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN