tico 0.1.0.dev250810__py3-none-any.whl → 0.1.0.dev250812__py3-none-any.whl

tico/__init__.py

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

tico/serialize/operators/adapters/__init__.py

@@ -0,0 +1 @@
+# DO NOT REMOVE THIS FILE

tico/serialize/operators/adapters/llama_rmsnorm.py

@@ -0,0 +1,35 @@
+# 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 contextlib import contextmanager
+
+import torch
+
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+
+
+def llama_rmsnorm_forward_adapter(self: LlamaRMSNorm, hidden_states: torch.Tensor):
+    return torch.ops.circle_custom.rms_norm(
+        hidden_states, self.weight, self.variance_epsilon
+    )
+
+
+@contextmanager
+def patched_llama_rmsnorm():
+    orig = LlamaRMSNorm.forward
+    LlamaRMSNorm.forward = llama_rmsnorm_forward_adapter
+    try:
+        yield
+    finally:
+        LlamaRMSNorm.forward = orig

tico/serialize/operators/op_rmsnorm.py

@@ -0,0 +1,65 @@
+# 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 CircleRMSNormArgs
+
+
+@register_node_visitor
+class RMSNormVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.circle_custom.rms_norm.default,
+    ]
+
+    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 = CircleRMSNormArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        weight = args.weight
+        eps = args.eps
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RMS_NORM, self._op_codes
+        )
+
+        inputs = [input, weight]
+        outputs = [node]
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.RmsNormOptions
+        )
+        option = circle.RmsNormOptions.RmsNormOptionsT()
+        option.epsilon = eps
+
+        operator.builtinOptions = option
+
+        return operator

tico/utils/register_custom_op.py

@@ -703,6 +703,28 @@ def CircleQuantizeMX():
         return input_
 
 
+def CircleRMSNorm():
+    @custom_op("circle_custom::rms_norm", mutates_args=())
+    def rms_norm(
+        hidden_states: torch.Tensor,
+        weight: Optional[torch.Tensor] = None,
+        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: Optional[torch.Tensor] = None,
+        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 +737,4 @@ def RegisterOps():
     CircleAvgPool2D()
     CircleInstanceNorm()
     CircleQuantizeMX()
+    CircleRMSNorm()

tico/utils/validate_args_kwargs.py

@@ -171,6 +171,19 @@ class CatArgs:
     dim: int = 0
 
 
+@enforce_type
+@dataclass
+class CircleRMSNormArgs:
+    """
+    This is not aten ops but custom op for RMSNorm.
+    circle_custom.rms_norm(Tensor input, Tensor? weight=None, float? eps=None) -> Tensor
+    """
+
+    input: torch.fx.Node
+    weight: Optional[torch.fx.Node]
+    eps: Optional[float]
+
+
 @enforce_type
 @dataclass
 class ClampArgs:
@@ -931,6 +944,19 @@ class ResizeNearestNeighborArgs:
     size: List[int]
 
 
+@enforce_type
+@dataclass
+class RMSNormArgs:
+    """
+    rms_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor
+    """
+
+    input: torch.fx.Node
+    normalized_shape: List[int]
+    weight: Optional[torch.fx.Node]
+    eps: Optional[float]
+
+
 @enforce_type
 @dataclass
 class RoundArgs:

{tico-0.1.0.dev250810.dist-info → tico-0.1.0.dev250812.dist-info}/METADATA

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

{tico-0.1.0.dev250810.dist-info → tico-0.1.0.dev250812.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=H67V6wHv1jk6smTy2N_nzF9vFR0U8Rf-Vsc19oWjhpM,1883
+tico/__init__.py,sha256=QzZEOB1tLI8yPD6dzOftGhdVFbFkAintnJP1vv89AJQ,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
@@ -166,6 +166,7 @@ tico/serialize/operators/op_relu6.py,sha256=ZWqEolfAKjOdUC1ZCg0iuu4dBhkJRxVYR2tU
 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_rmsnorm.py,sha256=vkJgg2YtTY9pjceTLh6gTZ-MN3EltnlEyAP5gVc5SiU,2216
 tico/serialize/operators/op_round.py,sha256=pe6w_TB4xGLu0iPv4Qo0a0fIkY9DgCgXk5127TWt8pE,1837
 tico/serialize/operators/op_rsqrt.py,sha256=yl2vd8InjhLPbE0vHIrEera6DVXlY9dLgO7yZZCH3RI,1837
 tico/serialize/operators/op_scalar_tensor.py,sha256=vDWxi4hXwyDJJhvfMR_QrBInw_No3WeU_M4gtfZqmbo,1928
@@ -186,6 +187,8 @@ tico/serialize/operators/op_unsqueeze.py,sha256=ZHhfVXSWEiwb2VDYX5uhxbGQyzZjKT7C
 tico/serialize/operators/op_view.py,sha256=xxE-GvTJ1UpcHst5KXYz3qKY-eJQvXKKrSZiA2O7E40,2593
 tico/serialize/operators/op_where.py,sha256=doE81GSwygrPBm3JIfN9w7kKXxeIYKxgk0eoY22QIcg,2845
 tico/serialize/operators/utils.py,sha256=lXGpEJW1h8U_-gfc6EWjvvSiq3yJ9P-v1v3EMRT_pSk,2954
+tico/serialize/operators/adapters/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
+tico/serialize/operators/adapters/llama_rmsnorm.py,sha256=6t3dhfNpR03eIjsmhymF2JKd6lCf7PvInqMf77c_BOE,1139
 tico/utils/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/convert.py,sha256=GgZwZtiqFzTdszfUQO0vcX39lKjs97gYwZ-Tiw_4Bbo,13222
 tico/utils/define.py,sha256=Ypgp7YffM4pgPl4Zh6TmogSn1OxGBMRw_e09qYGflZk,1467
@@ -200,20 +203,20 @@ 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=QN-8D71G_WAX3QQQ5CIwbEfFJZTQ3CvL4wCMiVddua4,3894
-tico/utils/register_custom_op.py,sha256=3-Yl6iYmx1qQA2igNHt4hYhQhQMkdPb7gF50LIY8yvc,27350
+tico/utils/register_custom_op.py,sha256=n91UtmPedoqhkR8fBNRbk9Msq79pn9DHNHlt99l2s_w,28142
 tico/utils/serialize.py,sha256=mEuusEzi82WFsz3AkowgWwxSLeo50JDxyOj6yYDQhEI,1914
 tico/utils/signature.py,sha256=R2GV0alRpXEbZISqPKyxCUWbgDcsrQ2ovbVG3737IzA,9595
 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
-tico/utils/validate_args_kwargs.py,sha256=CRj_SXMUUn6onsl8XLAt-zPZCFxR4C0XOCoaad_ZD4I,26689
+tico/utils/validate_args_kwargs.py,sha256=yikeUbYfSg2378wagEMXDlJeSRv8HKI2oxpjWarolec,27268
 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.dev250810.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250810.dist-info/METADATA,sha256=N1kcg1vk8kn6bLJgRUSZMZalY-mTn46jqLvVj9NGvR4,8450
-tico-0.1.0.dev250810.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250810.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250810.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250810.dist-info/RECORD,,
+tico-0.1.0.dev250812.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250812.dist-info/METADATA,sha256=VYlCGN-A_iI52ctvoTcdzg3UrtQ2zhZRCxyunaCdiqU,8450
+tico-0.1.0.dev250812.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250812.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250812.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250812.dist-info/RECORD,,