tico 0.1.0.dev250727__py3-none-any.whl → 0.1.0.dev250729__py3-none-any.whl

tico/__init__.py

@@ -29,7 +29,7 @@ __all__ = [
 ]
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250727"
+__version__ = "0.1.0.dev250729"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py

@@ -376,6 +376,12 @@ def _relu_handler(node, logger):
         quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
         node.meta[QPARAM_KEY] = _u8_to_i16(node.meta[QPARAM_KEY])
         logger.debug(f"quantize_per_tensor.default is inserted after {node.name}.")
+    elif qparam_dtype(inp) == "uint8" and qparam_dtype(node) == "int16":
+        quantize = _insert_quantize_op_after(node)
+
+        quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
+        node.meta[QPARAM_KEY] = _i16_to_u8(node.meta[QPARAM_KEY])
+        logger.debug(f"quantize_per_tensor.default is inserted after {node.name}.")
     else:
         raise NotYetSupportedError("Unsupported dtype")
 

tico/serialize/circle_graph.py

@@ -27,6 +27,7 @@ from tico.serialize.circle_mapping import (
     extract_circle_shape,
     str_to_circle_dtype,
     to_circle_dtype,
+    to_circle_shape,
 )
 from tico.serialize.pack import pack_buffer
 from tico.serialize.quant_param import QPARAM_KEY, QuantParam
@@ -151,7 +152,7 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         self.name_to_node[tensor.name] = node
         assert node.meta.get("val") is not None
         tensor.type = extract_circle_dtype(node)
-        tensor.shape, tensor.shapeSignature = extract_circle_shape(node)
+        tensor.shape, tensor.shapeSignature = extract_circle_shape(node)  # type: ignore[assignment]
 
         if QPARAM_KEY in node.meta:
             tensor.quantization = to_circle_qparam(node.meta[QPARAM_KEY])
@@ -186,7 +187,7 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         torch_t = torch.as_tensor(data=data)
         torch_t_shape = list(torch_t.size())
         tensor.type = to_circle_dtype(torch_dtype=torch_t.dtype)
-        tensor.shape = torch_t_shape
+        tensor.shape, tensor.shapeSignature = to_circle_shape(torch_t_shape)
 
         buffer = circle.Buffer.BufferT()
         buffer.data = torch_t.flatten().cpu().numpy().view(np.uint8)  # type: ignore[assignment]

tico/serialize/circle_mapping.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Tuple, TYPE_CHECKING, Union
+from typing import List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
     import torch.fx
@@ -132,18 +132,73 @@ def extract_circle_shape(node: torch.fx.Node) -> Tuple[List[int], Optional[List[
     return to_circle_shape(extract_shape(node))
 
 
-def to_circle_shape(torch_shape: torch.Size) -> Tuple[List[int], Optional[List[int]]]:
-    shape: List[int] = list(torch_shape)
-    shapeSignature: Optional[List[int]] = None
+def to_circle_shape(
+    torch_shape: Union[
+        torch.Size, Sequence[int | torch.SymInt]
+    ],  # Sequence[int | torch.SymInt] is added for type covariance
+) -> Tuple[List[int], Optional[List[int]]]:
 
-    if any(isinstance(s, torch.SymInt) for s in shape):
-        shapeSignature = shape.copy()
-        for idx, s in enumerate(shape):
+    if any(isinstance(s, torch.SymInt) for s in torch_shape):
+        # Follow dynamic shape spec
+        shape = []
+        shape_signature = []
+        for s in torch_shape:
             if isinstance(s, torch.SymInt):
-                shape[idx] = 1
-                shapeSignature[idx] = -1
-
-    return shape, shapeSignature
+                shape.append(1)
+                shape_signature.append(-1)
+            elif isinstance(s, int):
+                shape.append(s)
+                shape_signature.append(s)
+            else:
+                raise RuntimeError(f"Unsupported shape {torch_shape}")
+        return shape, shape_signature
+    else:
+        # Follow static shape spec
+        shape = []
+        shape_signature = None
+        for s in torch_shape:
+            if isinstance(s, int):
+                shape.append(s)
+            else:
+                assert False, "Cannot reach here"
+        return shape, shape_signature
+
+
+def validate_circle_shape(shape: List[int], shape_signature: Optional[List[int]]):
+    """
+    Validate circle tensor shape and shape_signature.
+    @ref https://github.com/Samsung/TICO/issues/244
+    """
+    if shape_signature is not None:
+        if len(shape_signature) == 0:
+            raise ValueError(
+                "Invalid circle shape: shape_signature must not be an empty list. "
+                "For static shapes, use None instead of []."
+            )
+        if len(shape) != len(shape_signature):
+            raise ValueError(
+                f"Invalid circle shape: shape and shape_signature must have same length: {shape} {shape_signature}"
+            )
+        if not all(isinstance(s, int) for s in shape_signature):
+            raise ValueError(
+                f"circle tensor shape_signature must be all integer values. {shape_signature}"
+            )
+        for s, ss in zip(shape, shape_signature):
+            if ss == -1:
+                # dynamic shape dimension
+                if s != 1:
+                    raise ValueError(
+                        f"Invalid circle shape: {s} {ss} {shape} {shape_signature}"
+                    )
+            else:
+                # static shape dimension
+                if s != ss:
+                    raise ValueError(
+                        f"Invalid circle shape: {s} {ss} {shape} {shape_signature}"
+                    )
+
+    if not all(isinstance(s, int) for s in shape):
+        raise ValueError(f"circle tensor shape must be all integer values. {shape}")
 
 
 # Return stride of node

tico/serialize/circle_serializer.py

@@ -26,7 +26,7 @@ from tico.serialize.circle_graph import CircleModel, CircleSubgraph
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import get_node_visitors
 from tico.utils import logging
-from tico.utils.serialize import finalise_tensor_names
+from tico.utils.serialize import finalise_tensor_names, validate_tensor_shapes
 
 
 multiple_output_ops = [
@@ -104,8 +104,10 @@ def build_circle(ep: ExportedProgram) -> bytes:
             graph.add_operator(circle_op)
             logger.debug(f"call_function: {node.name} ({opcode}) Op exported.")
 
-    # Register subgraph
     finalise_tensor_names(graph)
+    validate_tensor_shapes(graph)
+
+    # Register subgraph
     model.subgraphs.append(graph)
 
     # Encode operator codes

tico/serialize/operators/op_avg_pool2d.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_circle_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
@@ -57,12 +61,7 @@ class AvgPool2DVisitor(NodeVisitor):
             return True
 
     def has_same_padding(self, args: AvgPool2dArgs) -> bool:
-        input_shape, input_shape_signature = extract_circle_shape(args.input)
-
-        if input_shape_signature is not None:
-            # TODO: support dynamic shapes
-            raise NotImplementedError("Dynamic shape is not supported yet")
-
+        input_shape: torch.Size = extract_shape(args.input)
         kernel_size = args.kernel_size
         stride = args.stride
         assert stride
@@ -142,11 +141,7 @@ class AvgPool2DVisitor(NodeVisitor):
                 ],
                 dtype=torch.int32,
             )
-            input_shape, input_shape_signature = extract_circle_shape(input)
-
-            if input_shape_signature is not None:
-                raise RuntimeError("Dynamic shape is not supported yet.")
-
+            input_shape = extract_shape(input)
             input_dtype: int = extract_circle_dtype(input)
             padded_input_shape = [
                 input_shape[0],
@@ -156,11 +151,13 @@ class AvgPool2DVisitor(NodeVisitor):
             ]
             padded_input_shape[1] += padding[0] * 2
             padded_input_shape[2] += padding[1] * 2
+
             # 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_signature=None,
+                shape=padded_cshape,
+                shape_signature=padded_cshape_signature,
                 dtype=input_dtype,
                 source_node=node,
             )

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_circle_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,17 +115,13 @@ class Conv2dVisitor(NodeVisitor):
 
         assert groups == 1, "Only support group 1 conv2d"
 
-        input_shape, input_shape_signature = extract_circle_shape(input_)
-        output_shape, _ = extract_circle_shape(node)
-        weight_shape, _ = extract_circle_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)
 
-        if input_shape_signature is not None:
-            # TODO: support dynamic shapes
-            raise NotImplementedError("Dynamic shape is not supported yet")
-
         pad_decision = identify_padding(padding, input_shape, output_shape, stride)
 
         conv_input: torch.fx.Node | circle.Tensor.TensorT = input_
@@ -136,18 +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_signature=None,
+                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

@@ -110,7 +110,7 @@ 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: List[int] = dst_tensor.shapeSignature
+        dst_shape_signature: Optional[List[int]] = dst_tensor.shapeSignature
 
         if dst_shape_signature is not None:
             # TODO: support dynamic shape
@@ -134,7 +134,7 @@ class CopyVisitor(NodeVisitor):
 
         src_tensor: circle.Tensor.TensorT = self.graph.get_tensor(src)
         src_shape: List[int] = src_tensor.shape
-        src_shape_signature: List[int] = src_tensor.shapeSignature
+        src_shape_signature: Optional[List[int]] = src_tensor.shapeSignature
 
         if src_shape_signature is not None:
             # TODO: support dynamic shape

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,7 +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: List[int] = input_tensor.shapeSignature
+            input_shape_signature: Optional[List[int]] = input_tensor.shapeSignature
             cast_op_index = get_op_index(
                 circle.BuiltinOperator.BuiltinOperator.CAST, self._op_codes
             )

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_circle_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,17 +119,13 @@ class DepthwiseConv2dVisitor(NodeVisitor):
         dilation = args.dilation
         groups = args.groups
 
-        input_shape, input_shape_signature = extract_circle_shape(input_)  # OHWI
-        output_shape, _ = extract_circle_shape(node)  # OHWI
-        weight_shape, _ = extract_circle_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, len(weight_shape)
 
-        if input_shape_signature is not None:
-            # TODO: support dynamic shapes
-            raise NotImplementedError("Dynamic shape is not supported yet")
-
         assert weight_shape[0] == 1
         assert weight_shape[3] == output_shape[3]
         assert input_shape[3] == groups
@@ -150,18 +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_signature=None,
+                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_max_pool2d_with_indices.py

@@ -24,8 +24,8 @@ from circle_schema import circle
 from tico.serialize.circle_graph import CircleSubgraph
 from tico.serialize.circle_mapping import (
     extract_circle_dtype,
-    extract_circle_shape,
     extract_shape,
+    to_circle_shape,
 )
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
@@ -92,13 +92,15 @@ class MaxPool2DWithIndicesVisitor(NodeVisitor):
                 ],
                 dtype=torch.int32,
             )
-            input_shape, input_shape_signature = extract_circle_shape(input)
-
-            if input_shape_signature is not None:
-                # TODO: support dynamic shape
-                raise NotImplementedError("Padding with dynamic shape is not supported")
 
+            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],
@@ -107,18 +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_signature=None,
+                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_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,7 +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: List[int] = node_tensor.shapeSignature
+        node_shape_signature: Optional[List[int]] = node_tensor.shapeSignature
         op_index = get_op_index(
             circle.BuiltinOperator.BuiltinOperator.CAST, self._op_codes
         )

tico/serialize/operators/op_repeat.py

@@ -21,7 +21,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_circle_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
@@ -51,10 +55,7 @@ class RepeatVisitor(NodeVisitor):
             elif r < 0:
                 raise InvalidArgumentError("Only support positive repeat value")
 
-        tensor_shape, tensor_shape_signature = extract_circle_shape(input)
-        if tensor_shape_signature is not None:
-            # TODO: support dynamic shape
-            raise NotYetSupportedError("Repeat does not support dynamic shape yet.")
+        tensor_shape = extract_shape(input)
         assert len(tensor_shape) <= len(repeats)
         if len(tensor_shape) != len(repeats):
             # TODO Support len(tensor_shape) < len(repeats)
@@ -73,12 +74,16 @@ class RepeatVisitor(NodeVisitor):
             if r > 1:
                 # Except last created concat, a tensor should be created.
                 if repeat_dim_cnt > 1:
-                    repeated_shape = list(tensor_shape)
+                    repeated_shape: List[int | torch.SymInt] = list(tensor_shape)
                     repeated_shape[idx] = repeated_shape[idx] * r
+
+                    repeated_cshape, repeated_cshape_signature = to_circle_shape(
+                        repeated_shape
+                    )
                     concat_output = self.graph.add_tensor_from_scratch(
                         prefix=f"{node.name}_concat_{idx}",
-                        shape=repeated_shape,
-                        shape_signature=None,  # TODO: support dynamic shape
+                        shape=repeated_cshape,
+                        shape_signature=repeated_cshape_signature,
                         dtype=tensor_dtype,
                         source_node=node,
                     )

tico/serialize/operators/op_transpose_conv.py

@@ -24,6 +24,7 @@ from tico.serialize.circle_mapping import (
     circle_legalize_dtype_to,
     extract_circle_dtype,
     extract_circle_shape,
+    to_circle_shape,
 )
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
@@ -87,10 +88,6 @@ class TransposeConvVisitor(NodeVisitor):
         assert len(output_shape) == 4, len(output_shape)
         assert len(weight_shape) == 4, len(weight_shape)
 
-        if input_shape_signature is not None:
-            # TODO: support dynamic shapes
-            raise NotImplementedError("Dynamic shape is not supported yet")
-
         pad_decision = identify_padding(padding, input_shape, output_shape, stride)
 
         conv_input: torch.fx.Node | circle.Tensor.TensorT = input_
@@ -105,18 +102,21 @@ class TransposeConvVisitor(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_signature=None,
+                shape=pad_output_cshape,
+                shape_signature=pad_output_cshape_signature,
                 dtype=extract_circle_dtype(input_),
                 qparam=input_qparam,
                 source_node=node,

tico/utils/padding.py

@@ -15,6 +15,8 @@
 from enum import IntEnum
 from typing import NamedTuple, Optional, Sequence, Tuple, Union
 
+import torch
+
 from tico.utils.errors import InvalidArgumentError
 
 
@@ -37,8 +39,8 @@ class ConvPaddingInfo(NamedTuple):
 
 def identify_padding(
     padding: PaddingValue,
-    input_shape: Sequence[int],
-    output_shape: Sequence[int],
+    input_shape: Sequence[int | torch.SymInt] | torch.Size,
+    output_shape: Sequence[int | torch.SymInt] | torch.Size,
     stride: Sequence[int],
 ) -> ConvPaddingInfo:
     """

tico/utils/record_input.py

@@ -0,0 +1,92 @@
+# 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.
+
+import copy
+
+import inspect
+from contextlib import contextmanager
+from typing import Callable, List, Optional
+
+import torch.nn as nn
+
+
+class RecordingInput:
+    r"""Context-manager that records the input values of model::forward()
+
+    Recording input is useful for preparing example input for torch.export
+
+    Args:
+        condition: lambda to provide the condition whether to record or not
+
+            For examples, if you want to capture only args["past_key_values"] is not None,
+            conditon = lambda args_dict: args_dict["past_key_value"] is not None
+
+        input_to_remove: list of arg names to remove
+
+            Sometimes you would like to remove some arg values to make exported graph tidy or correct
+            For example, "past_key_values" may be not None, but just an empty cache. Then,
+            input_to_remove = [ "past_key_values" ]; makes the life easy
+
+    Example::
+        >>> with RecordingInput(model, input_to_remove=input_to_remove) as rec:
+        ...     outputs = model.generate(
+        ...     **inputs,
+        ...     )
+        ...     captured_input = rec.captured_input
+        >>> circle_model = tico.convert(model, captured_input)
+    """
+
+    def __init__(
+        self,
+        module: nn.Module,
+        condition: Callable[[dict], bool] = lambda args_dict: True,
+        *,
+        input_to_remove: Optional[List[str]] = [],
+    ):
+        self.module = module
+        self.forward_org = module.forward
+        self.condition = condition
+        self.input_to_remove = input_to_remove
+        self.sig = inspect.signature(self.forward_org)
+        self.args_names = [
+            name
+            for name in self.sig.parameters.keys()
+            if name not in ("self", "kwargs")
+        ]
+        self.captured_input = None
+
+    def __enter__(self):
+        def capture_and_forward(*args, **kwargs):
+            bound = self.sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+            args_dict = dict(bound.arguments)
+
+            def populate_args(args_dict, input_to_remove):
+                for key in input_to_remove:
+                    args_dict.pop(key, None)
+                args_tuple = tuple(
+                    args_dict.get(name, None) for name in self.args_names
+                )
+                return copy.deepcopy(args_tuple)
+
+            if self.condition(args_dict) and self.captured_input is None:
+                self.captured_input = populate_args(args_dict, self.input_to_remove)
+
+            return self.forward_org(*args, **kwargs)
+
+        self.module.forward = capture_and_forward
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.module.forward = self.forward_org

tico/utils/serialize.py

@@ -14,6 +14,7 @@
 
 
 from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.circle_mapping import validate_circle_shape
 from tico.utils.graph import get_module_name_chain
 
 
@@ -37,3 +38,18 @@ def finalise_tensor_names(
     for tensor in graph.tensors:
         if tensor.name in graph.name_to_node:
             tensor.name = f"{get_module_name_chain(graph.name_to_node[tensor.name])}::{tensor.name}"
+
+
+def validate_tensor_shapes(
+    graph: CircleSubgraph,
+) -> None:
+    """
+    Let's validate all tensors' shapes against their shape signatures.
+    """
+    for tensor in graph.tensors:
+        try:
+            validate_circle_shape(tensor.shape, tensor.shapeSignature)
+        except Exception as e:
+            raise ValueError(
+                f"Tensor {tensor.name} has invalid shape ({tensor.shape}), shape_signature ({tensor.shapeSignature})"
+            ) from e

{tico-0.1.0.dev250727.dist-info → tico-0.1.0.dev250729.dist-info}/METADATA

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

{tico-0.1.0.dev250727.dist-info → tico-0.1.0.dev250729.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=37Am_n1sr3CrCWSfMslOBX7P-iHXyaZ9IjQWSLozypc,1883
+tico/__init__.py,sha256=quSJ4KKyh76eIG3MrN15qViNKdudexAIk-h6x_L-hRc,1883
 tico/pt2_to_circle.py,sha256=gu3MD4Iqc0zMZcCZ2IT8oGbyj21CTSbT3Rgd9s2B_9A,2767
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=q5xMqGxTUZs4mFqt5c7i_y9U00fYgdMGl9nUqIVMlCo,1248
@@ -51,7 +51,7 @@ tico/experimental/quantization/evaluation/executor/circle_executor.py,sha256=eCC
 tico/experimental/quantization/evaluation/executor/triv24_executor.py,sha256=sUoXl6oOO2arAKaNjOBg7HiQja145_Jv6qgY7XtR7A8,5159
 tico/experimental/quantization/passes/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/passes/fold_quant_ops.py,sha256=bRHYSeHdSTaz3261skIkK5Aso2Lbv7ql0zFI9ICmbDY,7028
-tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py,sha256=4xwOzUrnMH4MQfdOHBHHXQo1EyWSkqNNVfBg47J5UCg,14721
+tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py,sha256=AtfK9kDnWyIWyVlwD4a0EEx_-5rW5Hmo5DuKZ-HyXH0,15069
 tico/experimental/quantization/passes/propagate_qparam_backward.py,sha256=TGtyW0Z2qOTgVIasBdGRgbwH31YYd6ek7OvLTmCV614,3118
 tico/experimental/quantization/passes/propagate_qparam_forward.py,sha256=RhUHGCR2RpBO5KYkQ7Z8U5u7HEwDq2wdKHLKAJCi-5c,5138
 tico/experimental/quantization/passes/quantize_bias.py,sha256=T7YxJ70N0tSK0FF9VJZA5iP0sHdnnsX9GX4AT4JDFSk,4325
@@ -96,9 +96,9 @@ tico/passes/remove_redundant_to_copy.py,sha256=tKy4XKkO2l33fMxVPQ_iFkUeFvP15kbPv
 tico/passes/restore_linear.py,sha256=xGJdNb-1CrkOKS9BnLbcblkZc6P2vVjKIi-7lRcs7Bk,4111
 tico/passes/segment_index_select.py,sha256=VVCKNLtYRkr9n5lGnlzEuQsQ0WVxEYXGchFrDnB1C40,5189
 tico/serialize/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
-tico/serialize/circle_graph.py,sha256=GSrrv6UGs2-sH5wJ19M-JLNimC-ZDsPzAO_ZswlHEU4,11945
-tico/serialize/circle_mapping.py,sha256=FEOGme-ewmZyTl_V34ErF7CI-KnDb98xmK_OcOKl4PA,6417
-tico/serialize/circle_serializer.py,sha256=bVKxAex_Nmc4O9RUFWIB_KhAuubVbhhA19CS3x967y8,10891
+tico/serialize/circle_graph.py,sha256=qvyul_HULoz7B_6RFKQ8s9RjEvMgPq-ynMVkZe8aqE4,12034
+tico/serialize/circle_mapping.py,sha256=c__AIHPi23lPugNJFolgMAKrw8j7gEeMaUQ1LAMSFnY,8542
+tico/serialize/circle_serializer.py,sha256=BGK9tltKkoL1h4rcrJUgDJIGlHst7aF3cZAKJk_GPWc,10950
 tico/serialize/pack.py,sha256=5HZ9kX3x6C6CyT_FWS6FRmvx_P7Dx21orjUNQxJ2xlo,1297
 tico/serialize/quant_param.py,sha256=6nbGKdqwMI9Cx9BLXJ9A9JU4qb770S8vTM1vCZRX3Eo,1342
 tico/serialize/operators/__init__.py,sha256=LIvXsNnN4yUCS2CGNQ5XW8p8oXDTV_WHWuOEAw1t6WY,990
@@ -110,17 +110,17 @@ tico/serialize/operators/op_alias_copy.py,sha256=Xu9OiILbGf8oddh8yTqovvLfgVs8XYV
 tico/serialize/operators/op_any.py,sha256=wrTXFQ1TWl-2ET2NGXAXI1dzfDDJsYtA71pyj2numPE,4968
 tico/serialize/operators/op_arange_start_step.py,sha256=0T5lWwh3TfsFStmVv0v5qG03KENRDBmMix08RXQ4D-U,2132
 tico/serialize/operators/op_argmax.py,sha256=ARyGHlmWVmzwCct93V5x1-VyKqhxMOvV8GuM8yQWXdo,2290
-tico/serialize/operators/op_avg_pool2d.py,sha256=OlxEezOtqdmc76d88Q5zziBBoOaIwtsHyhHMYKXHHog,7930
+tico/serialize/operators/op_avg_pool2d.py,sha256=eZl232FqYQsn2jUN_XjombHq_lzp2hf_AKwCLbZBxh8,7720
 tico/serialize/operators/op_bmm.py,sha256=AELjHC9ISFPIzEEl5Kr1s4GSNLZElwZmVZJWkEyCEoA,2189
 tico/serialize/operators/op_cat.py,sha256=XDYOh0XAyrM0TlxVm6Sa0OFFGrKk7aSDcGXC-hYX4gs,2204
 tico/serialize/operators/op_clamp.py,sha256=RRQVrzayDfN3PioCVJqa_yYOtcYwb5HHwkMe4E_YPmE,4408
 tico/serialize/operators/op_clone.py,sha256=vzDYJ8TS3tc2BAyd_z8nt5VqT1inpymSseMEhd9dva0,2394
 tico/serialize/operators/op_constant_pad_nd.py,sha256=OpP4AP-d1IFcWZolNa-o9ZxzXJQkMdG9WQ66soX3s-E,2675
-tico/serialize/operators/op_conv2d.py,sha256=I9OgOr1cBGFkx1Q6eyLLvY5DmKdaieMSNmVAuZ8pBa0,6842
-tico/serialize/operators/op_copy.py,sha256=z0GCaRP4dyJdJ1ajQ45BKmBRJlirWUPqYfnrqBtbeUM,6921
+tico/serialize/operators/op_conv2d.py,sha256=1_vouWXaF51gDLYg8z5Zlup0Tecq_ggAzvguiHzFffw,6828
+tico/serialize/operators/op_copy.py,sha256=boXHfl0bcvdBVl0tpzPMA_KBonh80vVqv61N3H5-PRU,6941
 tico/serialize/operators/op_cos.py,sha256=N12bNyuTQIxRnD0eHRPdFVzRQPMy1NFM4iM8oQ4lYzw,2034
-tico/serialize/operators/op_cumsum.py,sha256=Y8eVvtX57uPA4NIwS-Vbfjoiz_rPpc_wmViTQtCoz6s,3928
-tico/serialize/operators/op_depthwise_conv2d.py,sha256=S2naBMWAoUL3nJiV7RVg97bUvuQdB5VldE_r--rX0hA,7297
+tico/serialize/operators/op_cumsum.py,sha256=px9ZGUDDsdWjrql8Z1FdXfF-7CJhditxyNz5QRZbLiM,3948
+tico/serialize/operators/op_depthwise_conv2d.py,sha256=U6_nX2V31Evm-HLN9b3RKIVg-m8jyD-Nw1GdePUPPjY,7284
 tico/serialize/operators/op_dequantize_per_channel.py,sha256=aPcVxjdgvfSFoLnv9NL-RxO5vZYj8ulqriMP5LHIWs0,3133
 tico/serialize/operators/op_dequantize_per_tensor.py,sha256=u9aK_Xle9rDN0EHLE0YrCTlXY4Q53Ch9Di4qmx7ynps,2304
 tico/serialize/operators/op_div.py,sha256=WjeM2Ux7TyGlSNx2aVC783JvcL0xnY6FBYo1Q_kdb5Q,2201
@@ -144,7 +144,7 @@ tico/serialize/operators/op_logical_and.py,sha256=WhQ8knuq32BO-WhAqkOgpcUStPkjoP
 tico/serialize/operators/op_logical_not.py,sha256=ugrVcRqR3IvUUaiRVW5cArCYJbzmkcXp88QM846jCww,2129
 tico/serialize/operators/op_lt.py,sha256=_vA7dWpV9wVBxB7JL9pLQT9BsV91NGQBq_0auAtHK5Y,2080
 tico/serialize/operators/op_max_dim.py,sha256=nS_TZl5uq4uv1LwgBD9Wddyac4atKqBiIWKIyeXse2s,2519
-tico/serialize/operators/op_max_pool2d_with_indices.py,sha256=2Q1upWdChLpq1x7IaoP0eRNgqTt7dH7yqlRjhEKlnjM,6034
+tico/serialize/operators/op_max_pool2d_with_indices.py,sha256=i4iKZ262ytDKUt7bG9MiXuoKn--bgi-HWG24U5lvPPc,5919
 tico/serialize/operators/op_maximum.py,sha256=JjBr6gWEnuakLuk1_feotTHfIIm3s5YqWmqhUMpSPI0,1873
 tico/serialize/operators/op_mean.py,sha256=rVQZOxCJkHFY4kQBAS1HVK0HkcqxgkSy6zvEDLX_WYQ,2267
 tico/serialize/operators/op_minimum.py,sha256=fASjQVcTPCin02umQwFPdq2ss-Ve7S5A33J3QmmQ_wQ,1873
@@ -153,13 +153,13 @@ tico/serialize/operators/op_mul.py,sha256=si_VdYNyFbULb50SnXHOINh0dZQ2PhRB6Fzl54
 tico/serialize/operators/op_ne.py,sha256=xa2WJL2tYksxw7fIJic_D9ltLEseyCII8HpR32Oq8Do,1900
 tico/serialize/operators/op_neg.py,sha256=fkI3ExyD3QF-qtxBcXqQutPNDbNL8g7lZYE7CyD2wLk,2046
 tico/serialize/operators/op_permute.py,sha256=5DfX3pfZ5FDNmrSqx3-hRwPA7vm36z7BfG-nuyyBTsM,2282
-tico/serialize/operators/op_pow.py,sha256=G_GpdddmlE3opaCJtcfkV0qAWMau5uBa1sEB46qrNdo,5553
+tico/serialize/operators/op_pow.py,sha256=a-Nyy_s8d9nCIEAb5DacB1quDVmDu1VOHyAkD75u7Ts,5573
 tico/serialize/operators/op_prelu.py,sha256=0ZybL5pNvBrRvQGy4M6gELrjiEXEsb2wBDdU8x4D75I,1874
 tico/serialize/operators/op_quantize_per_tensor.py,sha256=w-vYxSPnN2gtx-pEkkcMGU0ZjiwaS4y1sxy56pKEq3E,3004
 tico/serialize/operators/op_reciprocal.py,sha256=6b9_bxjg_0EvgAitSv1MgBi4PJSEgm-21s5qtWI1UR4,2394
 tico/serialize/operators/op_relu.py,sha256=WXCR_chwEUBqjFIQ_4E2avwk-Acy76pmX20rJQCBTQo,1832
 tico/serialize/operators/op_relu6.py,sha256=ZWqEolfAKjOdUC1ZCg0iuu4dBhkJRxVYR2tUzpbvKQM,1829
-tico/serialize/operators/op_repeat.py,sha256=q05RQOheLxkWgj63Czrp6tGpjGuB0900FNRKLxXdNJ4,4524
+tico/serialize/operators/op_repeat.py,sha256=VrRxD31pT3hRGH-5n6ia3PJBXh_u0GvIl1hZZYFrKTQ,4507
 tico/serialize/operators/op_reshape.py,sha256=6wErQpmDX9mAmfJRCTg_cg1uOdJZqHm8Nux8dNI53Vg,2559
 tico/serialize/operators/op_resize_nearest_neighbor.py,sha256=dXaAnZ5M_ko_tH-HolxNpHFXkDUQ8x45myskojP5XZE,2771
 tico/serialize/operators/op_round.py,sha256=pe6w_TB4xGLu0iPv4Qo0a0fIkY9DgCgXk5127TWt8pE,1837
@@ -177,7 +177,7 @@ tico/serialize/operators/op_sub.py,sha256=yZskQJF0ylXVk02Uid8djPNIWDJ-0uHJar4UYh
 tico/serialize/operators/op_sum.py,sha256=B5aSwQMhyoBe2JYdE5nVQ3QeVDSzL-yuZZujsG08OdQ,2294
 tico/serialize/operators/op_tanh.py,sha256=rs7FsbQeUQ7Ak8RoQV9ymNGXHXRObojfY_SiqJiyqdA,1846
 tico/serialize/operators/op_to_copy.py,sha256=a8T0uPMavMO_md1a-4_0dlvDHyZS_xew0qB6xjf69rI,3934
-tico/serialize/operators/op_transpose_conv.py,sha256=PgM0mSIeiCMk-V2jQ2L1JdF-v3mHwdsdY4nTEp2GRro,5796
+tico/serialize/operators/op_transpose_conv.py,sha256=9NLnWpitfQzSDF-iAgw2fBA3YHL5y2Y8DQipeo8OvYA,5826
 tico/serialize/operators/op_unsqueeze.py,sha256=ZHhfVXSWEiwb2VDYX5uhxbGQyzZjKT7CrbBpVGxVHBU,2310
 tico/serialize/operators/op_view.py,sha256=xxE-GvTJ1UpcHst5KXYz3qKY-eJQvXKKrSZiA2O7E40,2593
 tico/serialize/operators/op_where.py,sha256=doE81GSwygrPBm3JIfN9w7kKXxeIYKxgk0eoY22QIcg,2845
@@ -192,11 +192,12 @@ tico/utils/graph.py,sha256=jD6m58m5JmN9mPfaROA9CW3406iJxmnukke00AuwRqI,9131
 tico/utils/installed_packages.py,sha256=J0FTwnkCGs0MxRWoCMYAqiwH7Z0GWFDLV--x-IndSp4,1017
 tico/utils/logging.py,sha256=IlbBWscsaHidI0dNqro1HEXAbIcbkR3BD5ukLy2m95k,1286
 tico/utils/model.py,sha256=pPOIjD0qjQirLibiRxxfjOR6efimOcDAd9R-74eus-k,1282
-tico/utils/padding.py,sha256=0iEcS5G3gSFySPzBz1m5nxHYZ6MlLzG_KdAFJRauwIg,3279
+tico/utils/padding.py,sha256=qKke-dJeeLHiRaePjDS66txrGyiYuipLVQeqLYad8uk,3349
 tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
 tico/utils/pytree_utils.py,sha256=jrk3N6X6LiUnBCX_gM1K9nywbVAJBVnszlTAgeIeDUc,5219
+tico/utils/record_input.py,sha256=FBtV00WWcXMXmg-Ujgvci9HjOmRJC1cVzx_WRNIF4MI,3324
 tico/utils/register_custom_op.py,sha256=3-Yl6iYmx1qQA2igNHt4hYhQhQMkdPb7gF50LIY8yvc,27350
-tico/utils/serialize.py,sha256=cBtEUfi_SU_9_v0cq2CNikzn8GnzEz2RwRvUH2NkWu4,1378
+tico/utils/serialize.py,sha256=mEuusEzi82WFsz3AkowgWwxSLeo50JDxyOj6yYDQhEI,1914
 tico/utils/torch_compat.py,sha256=oc6PztVsXdHcQ3iaVR90wLLxrGaj6zFHWZ8K9rRS6q8,1795
 tico/utils/trace_decorators.py,sha256=ddLIiKQfSaQrxgF1kNpwjFTQnXENzeSfcr1kuAW4jGI,3221
 tico/utils/utils.py,sha256=A5p3iAAxRGDsZJh4ybp-Qo3MX3vk5RrmSY-R3rXqVeI,12976
@@ -205,9 +206,9 @@ tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
 tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
 tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
-tico-0.1.0.dev250727.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250727.dist-info/METADATA,sha256=gE7opbcJdg8CiBvfdcjjvCuyrfQvy-q4BaBt8meVP84,8430
-tico-0.1.0.dev250727.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250727.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250727.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250727.dist-info/RECORD,,
+tico-0.1.0.dev250729.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250729.dist-info/METADATA,sha256=ZcnGD8K56o04Pt172XYRpuh-DHPoFWDhUhEpaCsy23k,8430
+tico-0.1.0.dev250729.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250729.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250729.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250729.dist-info/RECORD,,