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

tico/utils/record_input.py

@@ -0,0 +1,102 @@
+# 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 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)
+
+        for param in self.sig.parameters.values():
+            if param.kind == inspect.Parameter.KEYWORD_ONLY:
+                raise ValueError(f"Keyword-only parameter not supported: {param.name}")
+            if param.kind == inspect.Parameter.VAR_POSITIONAL:
+                raise ValueError(
+                    f"Var positional parameter not supported: {param.name}"
+                )
+
+        # NOTE: the name `kwargs` is removed since `kwargs` is a dict, not arg itself.
+        # args in kwargs are kept via sig.bind(*args, **kwargs) in capture_and_forward.
+        self.args_names = [
+            name
+            for name, param in self.sig.parameters.items()
+            if param.kind != inspect.Parameter.VAR_KEYWORD and name != "self"
+        ]
+        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/register_custom_op.py

@@ -31,9 +31,11 @@ def CircleResizeNearestNeighbor():
         W_scale_factor = size[2] / W
         if H_scale_factor != W_scale_factor:
             raise RuntimeError("Scale factor of H and W should be same.")
-        return torch.nn.functional.interpolate(
-            input_, scale_factor=H_scale_factor, mode="nearest"
+        permuted = torch.permute(input_, [0, 3, 1, 2])
+        resized = torch.nn.functional.interpolate(
+            permuted, scale_factor=H_scale_factor, mode="nearest"
         )
+        return torch.permute(resized, [0, 2, 3, 1])
 
     @register_fake("circle_custom::resize_nearest_neighbor")
     def _(input_: torch.Tensor, size: List[int]):
@@ -631,7 +633,7 @@ def CircleInstanceNorm():
         bias: Optional[torch.Tensor] = None,
         running_mean: Optional[torch.Tensor] = None,
         running_var: Optional[torch.Tensor] = None,
-        use_input_stats: bool = False,
+        use_input_stats: bool = True,
         momentum: float = 0.1,
         eps: float = 1e-05,
         cudnn_enabled: bool = False,
@@ -639,7 +641,7 @@ def CircleInstanceNorm():
         NHWC_to_NCHW = [0, 3, 1, 2]
         NCHW_input = torch.ops.aten.permute.default(input_, NHWC_to_NCHW)
 
-        args = [NCHW_input, weight, bias, None, None, False, momentum, eps, False]
+        args = [NCHW_input, weight, bias, None, None, True, momentum, eps, False]
         NCHW_output = torch.ops.aten.instance_norm.default(*args)
         NCHW_to_NHWC = [0, 2, 3, 1]
         NHWC_output = torch.ops.aten.permute.default(NCHW_output, NCHW_to_NHWC)
@@ -703,6 +705,28 @@ def CircleQuantizeMX():
         return input_
 
 
+def CircleRMSNorm():
+    @custom_op("circle_custom::rms_norm", mutates_args=())
+    def rms_norm(
+        hidden_states: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-05,
+    ) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + eps)
+        return weight * hidden_states.to(input_dtype)
+
+    @register_fake("circle_custom::rms_norm")
+    def _(
+        hidden_states: torch.Tensor,
+        weight: torch.Tensor,
+        eps: float = 1e-05,
+    ) -> torch.Tensor:
+        return hidden_states.new_empty(hidden_states.size())
+
+
 # Add custom ops to the torch namespace
 def RegisterOps():
     CircleResizeNearestNeighbor()
@@ -715,3 +739,4 @@ def RegisterOps():
     CircleAvgPool2D()
     CircleInstanceNorm()
     CircleQuantizeMX()
+    CircleRMSNorm()

tico/utils/serialize.py

@@ -12,11 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
-
-import torch
 
 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
 
 
@@ -40,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/utils/signature.py

@@ -0,0 +1,247 @@
+# 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 Sequence
+
+import numpy as np
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_mapping import to_circle_shape
+from tico.utils.dtype import circle_dtype_to_torch_dtype
+from tico.utils.installed_packages import is_dynamic_cache_available
+
+
+def is_dynamic_cache_instance(value):
+    if is_dynamic_cache_available():
+        from transformers.cache_utils import DynamicCache
+
+        return isinstance(value, DynamicCache)
+    else:
+        return False
+
+
+def flatten_and_convert_kwargs(kwargs: dict) -> dict[str, torch.Tensor]:
+    result = {}  # type: ignore[var-annotated]
+    for k, v in kwargs.items():
+        if v is None:
+            continue
+        elif isinstance(v, (list, tuple)):
+            # 1. handle list
+            def unpack_recursive(name, value, store=None):
+                if store is None:
+                    store = {}
+
+                if isinstance(value, (tuple, list)):
+                    for i, v in enumerate(value):
+                        # recursive call. Append index to name and explore lower level
+                        unpack_recursive(f"{name}_{i}", v, store)
+                else:
+                    # base type (scalar etc.) directly stored
+                    store[name] = value
+
+                return store
+
+            unpack_recursive(k, v, result)
+        elif is_dynamic_cache_instance(v):
+            # 2. handle DynamicCache
+            for idx, cache_val in enumerate(v.key_cache):
+                result[f"{k}_key_cache_{idx}"] = cache_val
+
+            for idx, cache_val in enumerate(v.value_cache):
+                result[f"{k}_value_cache_{idx}"] = cache_val
+        else:
+            result[k] = v
+
+    # 3. Convert to tensors
+    for k, v in result.items():
+        result[k] = v if isinstance(v, torch.Tensor) else torch.tensor(v)
+
+    return result
+
+
+def flatten_and_convert_args(args: Sequence) -> tuple:
+    result = []  # type: ignore[var-annotated]
+    for item in args:
+        if item is None:
+            continue
+
+        # 1. recursion on list and tuple
+        if isinstance(item, (list, tuple)):
+            result.extend(flatten_and_convert_args(item))
+            continue
+
+        # 2. handle DynamicCache
+        if is_dynamic_cache_available():
+            from transformers.cache_utils import DynamicCache
+
+            if isinstance(item, DynamicCache):
+                # NOTE The tensor order is: key_in → key_out → value_in → value_out
+                #
+                # Refer to https://github.com/huggingface/transformers/blob/3457e8e73e4f5532cc69059682b1ba4484d7e7e8/src/transformers/cache_utils.py#L557
+                # ```
+                # self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
+                # self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
+                # ```
+                result.extend(item.key_cache)
+                result.extend(item.value_cache)
+                continue
+
+        # 3. Convert to tensors
+        result.append(item if isinstance(item, torch.Tensor) else torch.tensor(item))
+
+    return tuple(result)
+
+
+class ModelInputSpec:
+    @classmethod
+    def load(cls, circle_path):
+        def load(circle_path: str) -> bytes:
+            with open(circle_path, "rb") as f:
+                buf = bytes(f.read())
+            return buf
+
+        circle_binary = load(circle_path)
+        return cls(circle_binary)
+
+    def __init__(self, circle_binary):
+        model = circle.Model.Model.GetRootAsModel(circle_binary, 0)
+        assert model.SubgraphsLength() == 1, "Only one subgraph is supported"
+
+        graph = model.Subgraphs(0)
+        tensors = [graph.Tensors(graph.Inputs(o)) for o in range(graph.InputsLength())]
+
+        self.names = [t.Name().decode("utf-8").split("::")[-1] for t in tensors]
+        self.shapes = [t.ShapeAsNumpy() for t in tensors]
+        self.shape_signatures = list(
+            map(
+                lambda x: None if (isinstance(x, int) and x == 0) else x,
+                (t.ShapeSignatureAsNumpy() for t in tensors),
+            )
+        )
+        self.types: list[torch.dtype] = [
+            circle_dtype_to_torch_dtype(t.Type()) for t in tensors
+        ]
+        self.name_to_idx = {name: idx for idx, name in enumerate(self.names)}
+
+    def bind(self, args, kwargs, check=True):
+        """Convert args and kwargs into an ordered list according to model input order"""
+        inputs = []
+        args = flatten_and_convert_args(args)
+        kwargs = flatten_and_convert_kwargs(kwargs)
+
+        arg_num = len(args) + len(kwargs)
+        m_input_num = len(self.names)
+        if arg_num != m_input_num:
+            raise ValueError(
+                f"Mismatch: number of model inputs and number of passed arguments are not the same: inputs({m_input_num}) != passed({arg_num}), input spec: {self.names}"
+            )
+
+        # 1. positional arguments
+        for i, val in enumerate(args):
+            name = self.names[i]
+            inputs.append(val)
+
+        # 2. keyword arguments
+        for idx in range(len(args), len(self.names)):
+            name = self.names[idx]
+            inputs.append(kwargs[name])
+
+        if check:
+            self.check_types(inputs)
+            self.check_shapes(inputs)
+
+        return inputs
+
+    def check_types(self, inputs):
+        """Check the types of input values"""
+        for i, (inp, ref_type) in enumerate(zip(inputs, self.types)):
+            # TODO: Support more data types (np array)
+            assert isinstance(
+                inp, (torch.Tensor | int | float)
+            ), f"Input '{self.names[i]}' type must be a torch tensor or scalar."
+
+            if isinstance(inp, torch.Tensor):
+                if inp.dtype != ref_type:
+                    raise TypeError(
+                        f"Input '{self.names[i]}' type {inp.dtype} != expected {ref_type}"
+                    )
+            else:
+                # Scalars (int, float)
+                if ref_type == torch.float32:
+                    if not isinstance(inp, (float)):
+                        raise TypeError(
+                            f"Input '{self.names[i]}' type {type(inp)} != expected {ref_type}"
+                        )
+                elif ref_type == torch.int64:
+                    if not isinstance(inp, (int)):
+                        raise TypeError(
+                            f"Input '{self.names[i]}' type {type(inp)} != expected {ref_type}"
+                        )
+                else:
+                    print(f"Unexpected ref_type: {ref_type}")
+
+    def check_shapes(self, inputs):
+        """Check the shapes of input values"""
+
+        def merge(shape, shape_sig):
+            """
+            Merge shape signature with shape
+            """
+            from copy import deepcopy
+
+            shape_merged = deepcopy(shape)
+            if shape_sig is not None:
+                for idx, ss in enumerate(shape_sig):
+                    if ss == -1:
+                        shape_merged[idx] = -1
+
+            return shape_merged
+
+        for i, (inp, ref_shape, ref_shape_sig) in enumerate(
+            zip(inputs, self.shapes, self.shape_signatures)
+        ):
+            # TODO: Support more data types (np array)
+            assert isinstance(
+                inp, (torch.Tensor | int | float)
+            ), f"Input '{self.names[i]}' type must be a torch tensor or scalar."
+
+            if isinstance(inp, torch.Tensor):  # Tensor
+                in_shape, in_shape_sig = to_circle_shape(inp.size())
+
+                if len(in_shape) != len(ref_shape):
+                    raise ValueError(
+                        f"Input '{self.names[i]}' has invalid rank {len(in_shape)}!= expected {len(ref_shape)}"
+                    )
+
+                in_merged_shape = merge(in_shape, in_shape_sig)
+                ref_merged_shape = merge(ref_shape, ref_shape_sig)
+                for in_shp, ref_shp in zip(in_merged_shape, ref_merged_shape):
+                    if ref_shp == -1:
+                        continue
+                    if in_shp == -1:
+                        raise ValueError(
+                            f"Input '{self.names[i]}' has unknown dimension {inp.size()} != expected shape({ref_shape}) / shape signature({ref_shape_sig}) "
+                        )
+                    if in_shp != ref_shp:
+                        raise ValueError(
+                            f"Input '{self.names[i]}' has wrong dimension {inp.size()} != expected shape({ref_shape}) / shape signature({ref_shape_sig}) "
+                        )
+            elif isinstance(inp, (int, float)):  # Scalar
+                if len(ref_shape) > 0:
+                    raise ValueError(
+                        f"Input '{self.names[i]}' has invalid rank {len(ref_shape)}"
+                    )
+            else:
+                print(f"Unexpected input type: {type(inp)}")

tico/utils/torch_compat.py

@@ -0,0 +1,52 @@
+# 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.
+
+"""
+Runtime **capability-detection helpers** for the `torch.export` stack.
+
+Instead of sprinkling version checks like `torch.__version__ >= "2.9"` throughout
+the codebase, import these helpers once and branch on the feature you need.
+
+Each probe executes only **once per process** thanks to `functools.lru_cache`,
+so the overhead is negligible.
+"""
+
+import functools
+
+import torch
+
+
+@functools.lru_cache(maxsize=None)
+def export_produces_slice() -> bool:
+    """
+    Compile a minimal model with `torch.export.export` and inspect its FX graph
+    to see whether an `aten.slice.Tensor` node appears.
+
+    Returns
+    -------
+    bool
+        * ``True``  — downstream passes should expect redundant **slice** nodes.
+        * ``False`` — downstream passes should expect only a **select** node.
+    """
+
+    class _Probe(torch.nn.Module):
+        def forward(self, x):  # simple slice: keep all dims except 3rd
+            return x[:, :, 1]
+
+        def get_example_inputs(self):
+            return (torch.randn(1, 4, 4),)
+
+    m = _Probe()
+    ep = torch.export.export(m, m.get_example_inputs())
+    return any(n.target == torch.ops.aten.slice.Tensor for n in ep.graph.nodes)

tico/utils/utils.py

@@ -21,7 +21,6 @@ from typing import List
 
 import torch
 from circle_schema import circle
-from packaging.version import Version
 from torch._guards import detect_fake_mode
 from torch.export import ExportedProgram
 from torch.utils import _pytree as pytree
@@ -29,10 +28,6 @@ from torch.utils import _pytree as pytree
 from tico.serialize.quant_param import QuantParam
 
 
-HAS_TORCH_OVER_25 = Version(torch.__version__) >= Version("2.5.0")
-HAS_TORCH_OVER_28_DEV = Version(torch.__version__) >= Version("2.8.0.dev")
-
-
 def get_fake_mode(exported_program: ExportedProgram):
     fake_mode = detect_fake_mode(
         tuple(
@@ -84,73 +79,70 @@ def enforce_type(callable):
     def check_types(*args, **kwargs):
         parameters = dict(zip(spec.args, args))
         parameters.update(kwargs)
-        for name, value in parameters.items():
-            if name == "self":
-                # skip 'self' in spec.args
-                continue
 
-            assert (
-                name in spec.annotations
-            ), f"All parameter require type hints. {name} needs a type hint"
-
-            type_hint = spec.annotations[name]
-
-            # Return tuple of flattened types.
-            # Q) What is flatten?
-            # A) Optional/Union is not included. Below are included.
-            # collections: List, Set, ...
-            # primitive types: int, str, ...
-            def _flatten_type(type_hint) -> tuple:
-                # `get_origin` maps Union[...] and Optional[...] varieties to Union
-                if typing.get_origin(type_hint) == typing.Union:
-                    # ex. typing.Union[list, int] -> (list, int)
-                    # ex. typing.Optional[torch.fx.Node] -> (torch.fx.Node, NoneType)
-                    actual_type = tuple(
-                        [_flatten_type(t) for t in typing.get_args(type_hint)]
-                    )
-                else:
-                    actual_type = (type_hint,)
-                return actual_type
+        # Return tuple of flattened types.
+        # Q) What is flatten?
+        # A) Optional/Union is not included. Below are included.
+        # collections: List, Set, ...
+        # primitive types: int, str, ...
+        def _flatten_type(type_hint) -> tuple:
+            # `get_origin` maps Union[...] and Optional[...] varieties to Union
+            if typing.get_origin(type_hint) == typing.Union:
+                # ex. typing.Union[list, int] -> (list, int)
+                # ex. typing.Optional[torch.fx.Node] -> (torch.fx.Node, NoneType)
+                actual_type = tuple(
+                    _flatten_type(t) for t in typing.get_args(type_hint)
+                )
+            else:
+                actual_type = (type_hint,)
+            return actual_type
 
-            type_hint = _flatten_type(type_hint)
+        # Return true if value matches with type_hint
+        # Return false otherwise
+        def _check_type(value, type_hint):
+            if type_hint == typing.Any:
+                return True
 
-            # Return true if value matches with type_hint
-            # Return false otherwise
-            def _check_type(value, type_hint):
-                if type_hint == typing.Any:
-                    return True
+            if isinstance(type_hint, tuple):
+                return any(_check_type(value, t) for t in type_hint)
 
-                if isinstance(type_hint, tuple):
-                    return any([_check_type(value, t) for t in type_hint])
+            if typing.get_origin(type_hint) in (list, set):
+                if not isinstance(value, typing.get_origin(type_hint)):
+                    return False
 
-                if typing.get_origin(type_hint) in (list, set):
-                    if not isinstance(value, typing.get_origin(type_hint)):
+                for v in value:
+                    if not any(_check_type(v, t) for t in typing.get_args(type_hint)):
                         return False
 
-                    for v in value:
-                        if not any(
-                            [_check_type(v, t) for t in typing.get_args(type_hint)]
-                        ):
-                            return False
+                return True
 
-                    return True
+            if typing.get_origin(type_hint) is dict:
+                if not isinstance(value, typing.get_origin(type_hint)):
+                    return False
 
-                if typing.get_origin(type_hint) == dict:
-                    if not isinstance(value, typing.get_origin(type_hint)):
+                for k, v in value.items():
+                    k_type, v_type = typing.get_args(type_hint)
+                    if not _check_type(k, k_type):
+                        return False
+                    if not _check_type(v, v_type):
                         return False
 
-                    for k, v in value.items():
-                        k_type, v_type = typing.get_args(type_hint)
-                        if not _check_type(k, k_type):
-                            return False
-                        if not _check_type(v, v_type):
-                            return False
+                return True
 
-                    return True
+            # TODO: Support more type hints
+            return isinstance(value, type_hint)
 
-                # TODO: Support more type hints
-                return isinstance(value, type_hint)
+        for name, value in parameters.items():
+            if name == "self":
+                # skip 'self' in spec.args
+                continue
 
+            assert (
+                name in spec.annotations
+            ), f"All parameter require type hints. {name} needs a type hint"
+
+            type_hint = spec.annotations[name]
+            type_hint = _flatten_type(type_hint)
             type_check_result = _check_type(value, type_hint)
             if not type_check_result:
                 raise ArgTypeError(