tico 0.1.0.dev250610__py3-none-any.whl → 0.1.0.dev250615__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,25 +55,22 @@ 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
+            permute = reshape_back_args.input
 
-            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 = (
                 reshape_front_args.input,
-                reshape_front_args.size,
+                reshape_front_args.shape,
             )
 
             # ---- condition: only leading unsqueeze ------------------------

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,17 +67,15 @@ 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
             reshape1_args = ReshapeArgs(*reshape1.args, **reshape1.kwargs)  # type: ignore[arg-type]
-            reshape1_input, reshape1_size = reshape1_args.input, reshape1_args.size
+            reshape1_input = reshape1_args.input
             # `(AxBxC) - aten.reshape` - (1xAxBxC)
             if [1] + list(extract_shape(reshape1_input)) != list(
                 extract_shape(reshape1)
@@ -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,19 +96,19 @@ 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
             reshape2_args = ReshapeArgs(*reshape2.args, **reshape2.kwargs)  # type: ignore[arg-type]
-            reshape2_input, reshape2_size = reshape2_args.input, reshape2_args.size
+            reshape2_input = reshape2_args.input
             # (1xAxCxB) - `aten.reshape - (AxCxB)
             if list(extract_shape(reshape2_input)) != [1] + list(
                 extract_shape(reshape2)
@@ -153,16 +151,13 @@ 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
             reshape1_args = ReshapeArgs(*reshape1.args, **reshape1.kwargs)  # type: ignore[arg-type]
-            reshape1_input, reshape1_size = reshape1_args.input, reshape1_args.size
+            reshape1_input = reshape1_args.input
             # `(AxBxC) - aten.reshape` - (1xAxBxC)
             if [1] + list(extract_shape(reshape1_input)) != list(
                 extract_shape(reshape1)
@@ -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,12 +178,12 @@ 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
             reshape2_args = ReshapeArgs(*reshape2.args, **reshape2.kwargs)  # type: ignore[arg-type]
-            reshape2_input, reshape2_size = reshape2_args.input, reshape2_args.size
+            reshape2_input, reshape2_size = reshape2_args.input, reshape2_args.shape
             # (Bx1xAxC) - `aten.reshape - (Bx(A*C))
             reshape2_input_shape = list(extract_shape(reshape2_input))
             assert len(reshape2_input_shape) == 4
@@ -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]
@@ -281,7 +270,7 @@ class RemoveRedundantReshapePattern3(PassBase):
             if not reshape_2.target in ops.aten.reshape:
                 continue
             reshape_2_args = ReshapeArgs(*reshape_2.args, **reshape_2.kwargs)  # type: ignore[arg-type]
-            reshape_2_input, reshape_2_size = reshape_2_args.input, reshape_2_args.size
+            reshape_2_input = reshape_2_args.input
             assert isinstance(reshape_2_input, torch.fx.Node), type(reshape_2_input)
             # reshape_3
             if not reshape_3.op == "call_function":
@@ -289,7 +278,7 @@ class RemoveRedundantReshapePattern3(PassBase):
             if not reshape_3.target in ops.aten.reshape:
                 continue
             reshape_3_args = ReshapeArgs(*reshape_3.args, **reshape_3.kwargs)  # type: ignore[arg-type]
-            reshape_3_input, reshape_3_size = reshape_3_args.input, reshape_3_args.size
+            reshape_3_input = reshape_3_args.input
             assert isinstance(reshape_3_input, torch.fx.Node), type(reshape_3_input)
 
             # Check condition
@@ -354,12 +343,11 @@ 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
+            reshape1_input, size = reshape1_args.input, reshape1_args.shape
             assert isinstance(reshape1_input, torch.fx.Node), type(reshape1_input)
             assert isinstance(size, list), type(size)
             for s in size:
@@ -370,13 +358,11 @@ 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_input, reshape2_size = reshape2_args.input, reshape2_args.size
+            reshape2_args = ReshapeArgs(*reshape2.args, **reshape2.kwargs)  # type: ignore[arg-type]
+            reshape2_input, reshape2_size = reshape2_args.input, reshape2_args.shape
             assert isinstance(reshape2_input, torch.fx.Node), type(reshape2_input)
             assert isinstance(reshape2_size, list), type(reshape2_size)
             for s in reshape2_size:
@@ -420,14 +406,11 @@ 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]
-            output_shape = args.size
+            output_shape = args.shape
             input_shape = list(extract_shape(args.input))
 
             if output_shape != input_shape:

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_reshape.py

@@ -48,7 +48,7 @@ class ReshapeVisitor(NodeVisitor):
         )
         args = ReshapeArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input = args.input
-        size = args.size
+        size = args.shape
 
         if isinstance(size, int):
             raise NotYetSupportedError("scalar size conversion is not supported yet.")