tico 0.1.0.dev250714__py3-none-any.whl → 0.1.0.dev251102__py3-none-any.whl

tico/serialize/operators/op_constant_pad_nd.py

@@ -28,6 +28,42 @@ from tico.utils.errors import InvalidArgumentError
 from tico.utils.validate_args_kwargs import ConstantPadNdArgs
 
 
+def convert_to_circle_padding(pad, input_shape_len):
+    MAX_RANK = 4
+
+    if not (1 <= input_shape_len <= MAX_RANK):
+        raise InvalidArgumentError(
+            f"Input rank must be between 1 and {MAX_RANK}, got {input_shape_len}"
+        )
+
+    if len(pad) % 2 != 0 or len(pad) < 2 or len(pad) > 8:
+        raise InvalidArgumentError(
+            f"Pad length must be an even number between 2 and 8, got {len(pad)}"
+        )
+
+    if len(pad) == 2:
+        padding = [[pad[0], pad[1]]]
+    elif len(pad) == 4:
+        padding = [[pad[2], pad[3]], [pad[0], pad[1]]]
+    elif len(pad) == 6:
+        padding = [[pad[4], pad[5]], [pad[2], pad[3]], [pad[0], pad[1]]]
+    elif len(pad) == 8:
+        padding = [
+            [pad[6], pad[7]],
+            [pad[4], pad[5]],
+            [pad[2], pad[3]],
+            [pad[0], pad[1]],
+        ]
+    else:
+        assert False, "Cannot reach here"
+
+    # Fill [0, 0] padding for the rest of dimension
+    while len(padding) < input_shape_len:
+        padding.insert(0, [0, 0])
+
+    return padding
+
+
 @register_node_visitor
 class ConstantPadNdVisitor(NodeVisitor):
     target: List[torch._ops.OpOverload] = [torch.ops.aten.constant_pad_nd.default]
@@ -45,19 +81,13 @@ class ConstantPadNdVisitor(NodeVisitor):
         val = args.value
 
         if val != 0:
-            raise InvalidArgumentError("Only support 0 value padding.")
+            raise InvalidArgumentError(f"Only support 0 value padding. pad:{pad}")
 
         input_shape_len = len(extract_shape(input_))
-        padding_size = [[pad[2], pad[3]], [pad[0], pad[1]]]
-        if input_shape_len == 3:
-            padding_size = [[0, 0]] + padding_size
-        elif input_shape_len == 4:
-            padding_size = [[0, 0], [0, 0]] + padding_size
-        else:
-            raise InvalidArgumentError("Only support 3D/4D inputs.")
-
-        paddings = torch.tensor(padding_size, dtype=torch.int32)
-        inputs = [input_, paddings]
+
+        padding = convert_to_circle_padding(pad, input_shape_len)
+
+        inputs = [input_, torch.tensor(padding, dtype=torch.int32)]
         outputs = [node]
 
         op_index = get_op_index(

tico/serialize/operators/op_conv2d.py

@@ -20,7 +20,11 @@ if TYPE_CHECKING:
 import torch
 from circle_schema import circle
 
-from tico.serialize.circle_mapping import extract_circle_dtype, extract_shape
+from tico.serialize.circle_mapping import (
+    extract_circle_dtype,
+    extract_shape,
+    to_circle_shape,
+)
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
@@ -111,9 +115,9 @@ class Conv2dVisitor(NodeVisitor):
 
         assert groups == 1, "Only support group 1 conv2d"
 
-        input_shape = list(extract_shape(input_))
-        output_shape = list(extract_shape(node))
-        weight_shape = list(extract_shape(weight))
+        input_shape = extract_shape(input_)
+        output_shape = extract_shape(node)
+        weight_shape = extract_shape(weight)
         assert len(input_shape) == 4, len(input_shape)
         assert len(output_shape) == 4, len(output_shape)
         assert len(weight_shape) == 4, len(weight_shape)
@@ -132,17 +136,21 @@ class Conv2dVisitor(NodeVisitor):
                 ],
                 dtype=torch.int32,
             )
-            pad_output_shape = [
+            pad_output_shape: List[int | torch.SymInt] = [
                 input_shape[0],
                 input_shape[1] + pad_h * 2,
                 input_shape[2] + pad_w * 2,
                 input_shape[3],
             ]
+            pad_output_cshape, pad_output_cshape_signature = to_circle_shape(
+                pad_output_shape
+            )
             # create padded output tensor
             input_qparam: Optional[QuantParam] = input_.meta.get(QPARAM_KEY)
             pad_output = self.graph.add_tensor_from_scratch(
                 prefix=f"{node.name}_input_pad_output",
-                shape=pad_output_shape,
+                shape=pad_output_cshape,
+                shape_signature=pad_output_cshape_signature,
                 dtype=extract_circle_dtype(input_),
                 qparam=input_qparam,
                 source_node=node,

tico/serialize/operators/op_copy.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, TYPE_CHECKING, Union
+from typing import Dict, List, Optional, TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
     import torch._ops
@@ -52,7 +52,15 @@ class CopyVisitor(NodeVisitor):
     def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
         super().__init__(op_codes, graph)
 
-    def check_to_do_broadcast(self, dst: List[int], src: List[int]) -> bool:
+    def check_to_do_broadcast(
+        self,
+        dst: List[int],
+        dst_sig: Optional[List[int]],
+        src: List[int],
+        src_sig: Optional[List[int]],
+    ) -> bool:
+        assert dst_sig is None
+        assert src_sig is None
         return dst != src
 
     def define_broadcast_to_node(
@@ -102,6 +110,12 @@ class CopyVisitor(NodeVisitor):
         # To connect 'dst' to Reshape node in the graph, 'dst' must be converted to Shape op.
         dst_tensor: circle.Tensor.TensorT = self.graph.get_tensor(dst)
         dst_shape: List[int] = dst_tensor.shape
+        dst_shape_signature: Optional[List[int]] = dst_tensor.shapeSignature
+
+        if dst_shape_signature is not None:
+            # TODO: support dynamic shape
+            raise NotYetSupportedError("Dynamic shape is not supported yet.")
+
         dst_shape_tensor = torch.as_tensor(dst_shape, dtype=torch.int32)
 
         dst_shape_shape = [len(dst_shape)]
@@ -110,6 +124,7 @@ class CopyVisitor(NodeVisitor):
         shape_output = self.graph.add_tensor_from_scratch(
             prefix=f"{dst_name}_shape_output",
             shape=dst_shape_shape,
+            shape_signature=None,
             dtype=circle.TensorType.TensorType.INT32,
             source_node=node,
         )
@@ -119,9 +134,16 @@ class CopyVisitor(NodeVisitor):
 
         src_tensor: circle.Tensor.TensorT = self.graph.get_tensor(src)
         src_shape: List[int] = src_tensor.shape
+        src_shape_signature: Optional[List[int]] = src_tensor.shapeSignature
+
+        if src_shape_signature is not None:
+            # TODO: support dynamic shape
+            raise NotYetSupportedError("Dynamic shape is not supported yet.")
 
         # The src tensor must be broadcastable with the dst tensor.
-        do_broadcast = self.check_to_do_broadcast(dst_shape, src_shape)
+        do_broadcast = self.check_to_do_broadcast(
+            dst_shape, dst_shape_signature, src_shape, src_shape_signature
+        )
         if do_broadcast:
             # create braodcastTo output tensor
             src_name: str = src.name
@@ -131,6 +153,7 @@ class CopyVisitor(NodeVisitor):
                 self.graph.add_tensor_from_scratch(
                     prefix=f"{src_name}_broadcast_to_output",
                     shape=dst_shape,
+                    shape_signature=dst_shape_signature,
                     dtype=src_type,
                     source_node=node,
                 )

tico/serialize/operators/op_cumsum.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, TYPE_CHECKING
+from typing import Dict, List, Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch._ops
@@ -57,6 +57,7 @@ class CumsumVisitor(NodeVisitor):
         if input_dtype == torch.int32:
             input_tensor: circle.Tensor.TensorT = self.graph.get_tensor(input)
             input_shape: List[int] = input_tensor.shape
+            input_shape_signature: Optional[List[int]] = input_tensor.shapeSignature
             cast_op_index = get_op_index(
                 circle.BuiltinOperator.BuiltinOperator.CAST, self._op_codes
             )
@@ -66,6 +67,7 @@ class CumsumVisitor(NodeVisitor):
                 prefix=cast_name,
                 dtype=cast_dtype,
                 shape=input_shape,
+                shape_signature=input_shape_signature,
                 source_node=node,
             )
             cast_operator = create_builtin_operator(

tico/serialize/operators/op_depthwise_conv2d.py

@@ -20,7 +20,11 @@ if TYPE_CHECKING:
 import torch
 from circle_schema import circle
 
-from tico.serialize.circle_mapping import extract_circle_dtype, extract_shape
+from tico.serialize.circle_mapping import (
+    extract_circle_dtype,
+    extract_shape,
+    to_circle_shape,
+)
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
@@ -115,12 +119,13 @@ class DepthwiseConv2dVisitor(NodeVisitor):
         dilation = args.dilation
         groups = args.groups
 
-        input_shape = list(extract_shape(input_))  # OHWI
-        output_shape = list(extract_shape(node))  # OHWI
-        weight_shape = list(extract_shape(weight))  # 1HWO
+        input_shape = extract_shape(input_)  # OHWI
+        output_shape = extract_shape(node)  # OHWI
+        weight_shape = extract_shape(weight)  # 1HWO
         assert len(input_shape) == 4, len(input_shape)
         assert len(output_shape) == 4, len(output_shape)
-        assert len(weight_shape) == 4
+        assert len(weight_shape) == 4, len(weight_shape)
+
         assert weight_shape[0] == 1
         assert weight_shape[3] == output_shape[3]
         assert input_shape[3] == groups
@@ -145,17 +150,22 @@ class DepthwiseConv2dVisitor(NodeVisitor):
                 ],
                 dtype=torch.int32,
             )
-            pad_output_shape = [
+            pad_output_shape: List[int | torch.SymInt] = [
                 input_shape[0],
                 input_shape[1] + pad_h * 2,
                 input_shape[2] + pad_w * 2,
                 input_shape[3],
             ]
+
+            pad_output_cshape, pad_output_cshape_signature = to_circle_shape(
+                pad_output_shape
+            )
             # create padded output tensor
             input_qparam: Optional[QuantParam] = input_.meta.get(QPARAM_KEY)
             pad_output = self.graph.add_tensor_from_scratch(
                 prefix=f"{node.name}_input_pad_output",
-                shape=pad_output_shape,
+                shape=pad_output_cshape,
+                shape_signature=pad_output_cshape_signature,
                 dtype=extract_circle_dtype(input_),
                 qparam=input_qparam,
                 source_node=node,

tico/serialize/operators/op_full_like.py

@@ -21,10 +21,8 @@ import torch
 from circle_schema import circle
 
 from tico.serialize.circle_graph import CircleSubgraph
-from tico.serialize.circle_mapping import to_circle_dtype
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
-from tico.serialize.operators.utils import create_builtin_operator, get_op_index
 from tico.utils.validate_args_kwargs import FullLikeArgs
 
 

tico/serialize/operators/op_index_select.py

@@ -49,7 +49,14 @@ class IndexSelectVisitor(NodeVisitor):
             self._op_codes,
         )
 
+        # TODO: Revise this to be simple
         dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
+        assert (
+            dim_i32.dim() == 0 or len(dim_i32) == 1
+        ), f"dim should be scalar: {dim_i32}"
+        dim_i32_item = dim_i32.item()
+        assert isinstance(dim_i32_item, int)
+
         inputs = [input, index]
         outputs = [node]
 
@@ -57,7 +64,7 @@ class IndexSelectVisitor(NodeVisitor):
 
         operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.GatherOptions
         option = circle.GatherOptions.GatherOptionsT()
-        option.axis = dim_i32
+        option.axis = dim_i32_item
 
         operator.builtinOptions = option
 

tico/serialize/operators/op_instance_norm.py

@@ -73,12 +73,6 @@ class InstanceNormVisitor(NodeVisitor):
         eps = args.eps
 
         # Ignore training-related args
-        running_mean = args.running_mean
-        running_var = args.running_var
-        use_input_stats = args.use_input_stats
-        momentum = args.momentum
-        cudnn_enabled = args.cudnn_enabled
-
         input_shape = list(extract_shape(input))
         assert len(input_shape) == 4, len(input_shape)
 

tico/serialize/operators/op_le.py

@@ -0,0 +1,54 @@
+# 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 Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import LeArgs
+
+
+@register_node_visitor
+class LeVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.le.Scalar,
+        torch.ops.aten.le.Tensor,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(self, node: torch.fx.Node) -> circle.Operator.OperatorT:
+        args = LeArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        other = args.other
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.LESS_EQUAL, self._op_codes
+        )
+
+        inputs = [input, other]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        return operator

tico/serialize/operators/op_log1p.py

@@ -23,7 +23,7 @@ from circle_schema import circle
 from tico.serialize.circle_graph import CircleSubgraph
 from tico.serialize.circle_mapping import (
     extract_circle_dtype,
-    extract_shape,
+    extract_circle_shape,
     extract_torch_dtype,
 )
 from tico.serialize.operators.hashable_opcode import OpCode
@@ -62,11 +62,12 @@ class Log1pVisitor(NodeVisitor):
         args = Log1pArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input = args.input
 
-        input_shape = list(extract_shape(input))
+        input_shape, input_shape_signature = extract_circle_shape(input)
         dst_dtype_circle = extract_circle_dtype(input)
         add_tensor: circle.Tensor.TensorT = self.graph.add_tensor_from_scratch(
             prefix=f"{input.name}_add",
             shape=input_shape,
+            shape_signature=input_shape_signature,
             dtype=dst_dtype_circle,
             source_node=node,
         )

tico/serialize/operators/op_max_pool2d_with_indices.py

@@ -22,7 +22,11 @@ import torch
 from circle_schema import circle
 
 from tico.serialize.circle_graph import CircleSubgraph
-from tico.serialize.circle_mapping import extract_circle_dtype, extract_shape
+from tico.serialize.circle_mapping import (
+    extract_circle_dtype,
+    extract_shape,
+    to_circle_shape,
+)
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import (
@@ -88,8 +92,15 @@ class MaxPool2DWithIndicesVisitor(NodeVisitor):
                 ],
                 dtype=torch.int32,
             )
-            input_shape = list(extract_shape(input))
+
+            input_shape = extract_shape(input)
             input_dtype: int = extract_circle_dtype(input)
+
+            input_qparam: Optional[QuantParam] = (
+                input.meta[QPARAM_KEY] if QPARAM_KEY in input.meta else None
+            )
+
+            # create padded input tensor
             padded_input_shape = [
                 input_shape[0],
                 input_shape[1],
@@ -98,17 +109,16 @@ class MaxPool2DWithIndicesVisitor(NodeVisitor):
             ]
             padded_input_shape[1] += padding[0] * 2
             padded_input_shape[2] += padding[1] * 2
-            input_qparam: Optional[QuantParam] = (
-                input.meta[QPARAM_KEY] if QPARAM_KEY in input.meta else None
-            )
-            # create padded input tensor
+            padded_cshape, padded_cshape_signature = to_circle_shape(padded_input_shape)
             padded_input_tensor = self.graph.add_tensor_from_scratch(
                 prefix=f"{input.name}_pad_output",
-                shape=padded_input_shape,
+                shape=padded_cshape,
+                shape_signature=padded_cshape_signature,
                 dtype=input_dtype,
                 qparam=input_qparam,
                 source_node=node,
             )
+
             if input_qparam is not None:
                 padding_value = get_integer_dtype_min(input_qparam.dtype)
             else:

tico/serialize/operators/op_mm.py

@@ -20,7 +20,7 @@ if TYPE_CHECKING:
 import torch
 from circle_schema import circle
 
-from tico.serialize.circle_graph import CircleSubgraph, is_const
+from tico.serialize.circle_graph import CircleSubgraph
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
@@ -28,9 +28,9 @@ from tico.utils.validate_args_kwargs import MatmulArgs
 
 
 @register_node_visitor
-class MatmulDefaultVisitor(NodeVisitor):
+class MatmulVisitor(NodeVisitor):
     """
-    Convert matmul to equavalent BatchMatMul or FullyConnected with Transpose.
+    Convert matmul to Circle BatchMatMul
     """
 
     target: List[torch._ops.OpOverload] = [torch.ops.aten.mm.default]
@@ -38,130 +38,7 @@ class MatmulDefaultVisitor(NodeVisitor):
     def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
         super().__init__(op_codes, graph)
 
-    # NOTE: Matmul is equivalent to Batch MatMul (batch=1)
-    def define_bmm_node(self, inputs, outputs) -> circle.Operator.OperatorT:
-        def set_bmm_option(operator):
-            operator.builtinOptionsType = (
-                circle.BuiltinOptions.BuiltinOptions.BatchMatMulOptions
-            )
-            option = circle.BatchMatMulOptions.BatchMatMulOptionsT()
-            option.adjointLhs, option.adjointRhs = False, False
-            option.asymmetricQuantizeInputs = False
-            operator.builtinOptions = option
-
-        op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.BATCH_MATMUL, self._op_codes
-        )
-        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
-        set_bmm_option(operator)
-
-        return operator
-
-    def define_transpose_node(self, inputs, outputs) -> circle.Operator.OperatorT:
-        def set_transpose_option(operator):
-            operator.builtinOptionsType = (
-                circle.BuiltinOptions.BuiltinOptions.TransposeOptions
-            )
-            option = circle.TransposeOptions.TransposeOptionsT()
-            operator.builtinOptions = option
-
-        transpose_op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.TRANSPOSE, self._op_codes
-        )
-        operator = create_builtin_operator(
-            self.graph, transpose_op_index, inputs, outputs
-        )
-        set_transpose_option(operator)
-        return operator
-
-    def define_fc_node(self, inputs, outputs) -> circle.Operator.OperatorT:
-        def set_fc_option(operator):
-            operator.builtinOptionsType = (
-                circle.BuiltinOptions.BuiltinOptions.FullyConnectedOptions
-            )
-            option = circle.FullyConnectedOptions.FullyConnectedOptionsT()
-
-            option.fusedActivationFunction = (
-                circle.ActivationFunctionType.ActivationFunctionType.NONE
-            )
-            option.weightsFormat = (
-                circle.FullyConnectedOptionsWeightsFormat.FullyConnectedOptionsWeightsFormat.DEFAULT
-            )
-            option.keepNumDims = False
-            option.asymmetricQuantizeInputs = False
-            option.quantizedBiasType = circle.TensorType.TensorType.FLOAT32
-
-            operator.builtinOptions = option
-
-        fc_op_index = get_op_index(
-            circle.BuiltinOperator.BuiltinOperator.FULLY_CONNECTED, self._op_codes
-        )
-        operator = create_builtin_operator(self.graph, fc_op_index, inputs, outputs)
-        set_fc_option(operator)
-        return operator
-
-    """
-    Define FullyConnnected with Tranpose operator.
-    Note that those sets of operators are equivalent.
-    (1) Matmul
-    matmul( lhs[H, K], rhs[K, W'] ) -> output(H, W')
-    
-    (2) Transpose + FullyConneccted
-    transpose( rhs[K, W'] ) -> trs_output[W', K]
-    fullyconnected( lhs[H, K], trs_output[W', K] ) -> output(H, W')
-    """
-
-    def define_fc_with_transpose(
-        self, node, inputs, outputs
-    ) -> circle.Operator.OperatorT:
-        lhs, rhs = inputs
-
-        # get transpose shape
-        rhs_tid: int = self.graph.get_tid_registered(rhs)
-        rhs_tensor: circle.Tensor.TensorT = self.graph.tensors[rhs_tid]
-        rhs_name: str = rhs.name
-        rhs_type: int = rhs_tensor.type
-        rhs_shape: List[int] = rhs_tensor.shape
-        assert len(rhs_shape) == 2, len(rhs_shape)
-        rhs_shape_transpose = [rhs_shape[1], rhs_shape[0]]
-
-        # create transpose output tensor
-        trs_output = self.graph.add_tensor_from_scratch(
-            prefix=f"{rhs_name}_transposed_output",
-            shape=rhs_shape_transpose,
-            dtype=rhs_type,
-            source_node=node,
-        )
-        trs_perm = self.graph.add_const_tensor(data=[1, 0], source_node=node)
-        trs_operator = self.define_transpose_node([rhs, trs_perm], [trs_output])
-        self.graph.add_operator(trs_operator)
-
-        # define fc node
-        fc_input = lhs
-        fc_weight = trs_output
-        fc_shape = [fc_weight.shape[0]]
-        fc_bias = self.graph.add_const_tensor(
-            data=[0.0] * fc_shape[0], source_node=node
-        )
-
-        operator = self.define_fc_node([fc_input, fc_weight, fc_bias], outputs)
-
-        return operator
-
-    def define_node(
-        self, node: torch.fx.Node, prior_latency=True
-    ) -> circle.Operator.OperatorT:
-        """
-        NOTE: Possibility of accuracy-latency trade-off
-        From ONE compiler's perspective:
-        - BMM uses per-tensor quantization for both rhs and lhs.
-        - FC uses per-channel quantization for weight and per-tensor for input.
-        Thus, FC is better in terms of accuracy.
-        FC necessarily involves an additional transpose operation to be identical with mm.
-        If transposed operand is const, it can be optimized by constant folding.
-        Thus, convert FC only if tranpose can be folded.
-        TODO set prior_latency outside
-        """
+    def define_node(self, node: torch.fx.Node) -> circle.Operator.OperatorT:
         args = MatmulArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input = args.input
         other = args.other
@@ -169,9 +46,16 @@ class MatmulDefaultVisitor(NodeVisitor):
         inputs = [input, other]
         outputs = [node]
 
-        if not is_const(other) and prior_latency:
-            operator = self.define_bmm_node(inputs, outputs)
-        else:
-            operator = self.define_fc_with_transpose(node, inputs, outputs)
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.BATCH_MATMUL, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.BatchMatMulOptions
+        )
+        option = circle.BatchMatMulOptions.BatchMatMulOptionsT()
+        option.adjointLhs, option.adjointRhs = False, False
+        option.asymmetricQuantizeInputs = False
+        operator.builtinOptions = option
 
         return operator

tico/serialize/operators/op_mul.py

@@ -66,10 +66,7 @@ class MulTensorVisitor(BaseMulVisitor):
         self,
         node: torch.fx.Node,
     ) -> circle.Operator.OperatorT:
-        args = MulTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-        input = args.input
-        other = args.other
-
+        _ = MulTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         operator = super().define_node(
             node,
         )
@@ -88,10 +85,7 @@ class MulScalarVisitor(BaseMulVisitor):
         self,
         node: torch.fx.Node,
     ) -> circle.Operator.OperatorT:
-        args = MulScalarArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-        input = args.input
-        other = args.other
-
+        _ = MulScalarArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         operator = super().define_node(
             node,
         )

tico/serialize/operators/op_pow.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, TYPE_CHECKING
+from typing import Dict, List, Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch._ops
@@ -36,6 +36,7 @@ class BasePowVisitor(NodeVisitor):
         assert isinstance(node, torch.fx.Node), type(node)
         node_tensor: circle.Tensor.TensorT = self.graph.get_tensor(node)
         node_shape: List[int] = node_tensor.shape
+        node_shape_signature: Optional[List[int]] = node_tensor.shapeSignature
         op_index = get_op_index(
             circle.BuiltinOperator.BuiltinOperator.CAST, self._op_codes
         )
@@ -45,6 +46,7 @@ class BasePowVisitor(NodeVisitor):
             prefix=cast_name,
             dtype=cast_dtype,
             shape=node_shape,
+            shape_signature=node_shape_signature,
             source_node=node,
         )
         cast_operator = create_builtin_operator(