tico 0.1.0.dev250824__py3-none-any.whl → 0.1.0.dev250825__py3-none-any.whl

tico/__init__.py

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

tico/experimental/quantization/ptq/wrappers/nn/__init__.py

@@ -1,5 +1,11 @@
+from tico.experimental.quantization.ptq.wrappers.nn.quant_layernorm import (
+    QuantLayerNorm,
+)
 from tico.experimental.quantization.ptq.wrappers.nn.quant_linear import QuantLinear
+from tico.experimental.quantization.ptq.wrappers.nn.quant_silu import QuantSiLU
 
 __all__ = [
+    "QuantLayerNorm",
     "QuantLinear",
+    "QuantSiLU",
 ]

tico/experimental/quantization/ptq/wrappers/nn/quant_layernorm.py

@@ -0,0 +1,184 @@
+# 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 Iterable, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from tico.experimental.quantization.ptq.mode import Mode
+from tico.experimental.quantization.ptq.quant_config import QuantConfig
+from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
+    QuantModuleBase,
+)
+from tico.experimental.quantization.ptq.wrappers.registry import register
+
+
+@register(nn.LayerNorm)
+class QuantLayerNorm(QuantModuleBase):
+    """
+    QuantLayerNorm — drop-in replacement for nn.LayerNorm that quantizes
+    the elementary steps:
+        1) μ = mean(x, dims)                (mean)
+        2) c = x - μ                        (sub)
+        3) s = c * c                        (square)
+        4) v = mean(s, dims)                (variance)
+        5) e = v + eps                      (add-eps)
+        6) r = rsqrt(e)                     (rsqrt)
+        7) n = c * r                        (normalize)
+        8) y = (n * γ) + β                  (affine), with:
+             • affine_mul : n * γ
+             • affine_add : (n * γ) + β
+    """
+
+    def __init__(
+        self,
+        fp: nn.LayerNorm,
+        *,
+        qcfg: Optional[QuantConfig] = None,
+        fp_name: Optional[str] = None
+    ):
+        super().__init__(qcfg, fp_name=fp_name)
+        self.module = fp
+        self.eps = torch.tensor(self.module.eps)
+        # Number of trailing dims participating in normalization
+        # (PyTorch stores normalized_shape as a tuple even if an int was passed)
+        self._norm_ndim: int = len(fp.normalized_shape)  # safe for int→tuple
+
+        # Activation / intermediate observers
+        self.act_in_obs = self._make_obs("act_in")
+        self.mean_obs = self._make_obs("mean")
+        self.centered_obs = self._make_obs("centered")
+        self.square_obs = self._make_obs("square")
+        self.var_obs = self._make_obs("var")
+        self.eps_obs = self._make_obs("eps")
+        self.add_eps_obs = self._make_obs("add_eps")
+        self.inv_std_obs = self._make_obs("inv_std")
+        self.norm_obs = self._make_obs("norm")
+        self.act_out_obs = self._make_obs("act_out")
+
+        # Optional affine parameter observers (γ, β)
+        self.weight_obs = None
+        self.bias_obs = None
+        self.affine_mul_obs = None
+        self.affine_add_obs = None
+        if self.module.elementwise_affine:
+            if self.module.weight is not None:
+                self.weight_obs = self._make_obs("weight")
+            if self.module.bias is not None:
+                self.bias_obs = self._make_obs("bias")
+            # Per-op observers for (n * w) and (+ b)
+            self.affine_mul_obs = self._make_obs("affine_mul")
+            self.affine_add_obs = self._make_obs("affine_add")
+
+    def enable_calibration(self) -> None:
+        """
+        Switch to CALIB mode and collect *fixed* ranges for affine params
+        immediately, since they do not change across inputs.
+        """
+        super().enable_calibration()
+        if self.module.elementwise_affine:
+            if self.weight_obs is not None and self.module.weight is not None:
+                self.weight_obs.collect(self.module.weight)
+            if self.bias_obs is not None and self.module.bias is not None:
+                self.bias_obs.collect(self.module.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Determine reduction dims (last self._norm_ndim axes)
+        # Example: if x.ndim=4 and norm_ndim=2 → dims=(2,3)
+        dims = tuple(range(x.dim() - self._norm_ndim, x.dim()))
+
+        # 0) input
+        x_q = self._fq(x, self.act_in_obs)
+
+        # 1) mean
+        mu = x_q.mean(dim=dims, keepdim=True)
+        mu_q = self._fq(mu, self.mean_obs)
+
+        # 2) center
+        c = x_q - mu_q
+        c_q = self._fq(c, self.centered_obs)
+
+        # 3) square (elementwise mul)
+        s = c_q * c_q
+        s_q = self._fq(s, self.square_obs)
+
+        # 4) variance (via squared mean)
+        v = s_q.mean(dim=dims, keepdim=True)
+        v_q = self._fq(v, self.var_obs)
+
+        # 5) add eps
+        eps_q = self._fq(self.eps, self.eps_obs)
+        e = v_q + eps_q
+        e_q = self._fq(e, self.add_eps_obs)
+
+        # 6) inverse std
+        r = torch.rsqrt(e_q)
+        r_q = self._fq(r, self.inv_std_obs)
+
+        # 7) normalize
+        n = c_q * r_q
+        n_q = self._fq(n, self.norm_obs)
+
+        # 8) optional affine
+        if self.module.elementwise_affine:
+            w = self.module.weight
+            b = self.module.bias
+            if self._mode is Mode.QUANT:
+                if self.weight_obs is not None and w is not None:
+                    w = self.weight_obs.fake_quant(w)  # type: ignore[assignment]
+                if self.bias_obs is not None and b is not None:
+                    b = self.bias_obs.fake_quant(b)  # type: ignore[assignment]
+            y = n_q
+            # 8a) n * w  (fake-quant the result of the mul)
+            if w is not None:
+                y = y * w
+                if self.affine_mul_obs is not None:
+                    y = self._fq(y, self.affine_mul_obs)
+
+            # 8b) (+ b)  (fake-quant the result of the add)
+            if b is not None:
+                y = y + b
+                if self.affine_add_obs is not None:
+                    y = self._fq(y, self.affine_add_obs)
+        else:
+            y = n_q
+
+        # 9) output activation
+        return self._fq(y, self.act_out_obs)
+
+    def _all_observers(self) -> Iterable:
+        obs: Tuple = (
+            self.act_in_obs,
+            self.mean_obs,
+            self.centered_obs,
+            self.square_obs,
+            self.var_obs,
+            self.eps_obs,
+            self.add_eps_obs,
+            self.inv_std_obs,
+            self.norm_obs,
+            self.act_out_obs,
+        )
+        # Insert affine param observers if present
+        if self.module.elementwise_affine:
+            if self.weight_obs is not None:
+                obs = (self.weight_obs,) + obs
+            if self.bias_obs is not None:
+                obs = obs + (self.bias_obs,)
+            if self.affine_mul_obs is not None:
+                obs = obs + (self.affine_mul_obs,)
+            if self.affine_add_obs is not None:
+                obs = obs + (self.affine_add_obs,)
+        return obs

tico/experimental/quantization/ptq/wrappers/nn/quant_silu.py

@@ -0,0 +1,61 @@
+# 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 Optional
+
+import torch
+import torch.nn as nn
+
+from tico.experimental.quantization.ptq.quant_config import QuantConfig
+from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
+    QuantModuleBase,
+)
+from tico.experimental.quantization.ptq.wrappers.registry import register
+
+
+@register(nn.SiLU)
+class QuantSiLU(QuantModuleBase):
+    """
+    QuantSiLU — drop-in replacement for nn.SiLU that quantizes
+    both intermediate tensors:
+        • s  = sigmoid(x)   (logistic)
+        • y  = x * s        (mul)
+    """
+
+    def __init__(
+        self,
+        fp: nn.SiLU,
+        *,
+        qcfg: Optional[QuantConfig] = None,
+        fp_name: Optional[str] = None
+    ):
+        super().__init__(qcfg, fp_name=fp_name)
+        self.act_in_obs = self._make_obs("act_in")
+        self.sig_obs = self._make_obs("sigmoid")
+        self.mul_obs = self._make_obs("mul")
+        self.module = fp
+
+    def forward(self, x: torch.Tensor):
+        x_q = self._fq(x, self.act_in_obs)
+
+        s = torch.sigmoid(x_q)
+        s = self._fq(s, self.sig_obs)
+
+        y = x * s
+        y = self._fq(y, self.mul_obs)
+
+        return y
+
+    def _all_observers(self):
+        return (self.act_in_obs, self.sig_obs, self.mul_obs)

tico/experimental/quantization/ptq/wrappers/quant_module_base.py

@@ -46,10 +46,24 @@ class QuantModuleBase(nn.Module, ABC):
         self.fp_name = fp_name
 
     def _child_quant_modules(self):
-        """Yield direct children that are QuantModuleBase."""
-        for m in self.children():
+        """
+        Yield immediate QuantModuleBase *descendants*, skipping over pure containers
+        (e.g., ModuleList/Sequential/ModuleDict). Once a QuantModuleBase is found,
+        do NOT descend into it here—let recursion happen level by level.
+        """
+        seen = set()
+        stack = list(self.children())  # start from direct children
+
+        while stack:
+            m = stack.pop()
             if isinstance(m, QuantModuleBase):
-                yield m
+                if id(m) not in seen:
+                    seen.add(id(m))
+                    yield m
+                # IMPORTANT: do not recurse into `m` here; its own call will handle its subtree
+            elif isinstance(m, (nn.ModuleList, nn.ModuleDict, nn.Sequential)):
+                # `m` is a container or a non-quant leaf: keep descending until we hit quant modules
+                stack.extend(list(m.children()))
 
     def enable_calibration(self) -> None:
         self._mode = Mode.CALIB

tico/experimental/quantization/ptq/wrappers/registry.py

@@ -24,7 +24,10 @@ from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
 _WRAPPERS: Dict[Type[nn.Module], Type[QuantModuleBase]] = {}
 _IMPORT_ONCE = False
 _CORE_MODULES = (
+    "tico.experimental.quantization.ptq.wrappers.quant_elementwise",
+    "tico.experimental.quantization.ptq.wrappers.nn.quant_layernorm",
     "tico.experimental.quantization.ptq.wrappers.nn.quant_linear",
+    "tico.experimental.quantization.ptq.wrappers.nn.quant_silu",
     # add future core wrappers here
 )
 

{tico-0.1.0.dev250824.dist-info → tico-0.1.0.dev250825.dist-info}/METADATA

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

{tico-0.1.0.dev250824.dist-info → tico-0.1.0.dev250825.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=KpoO-uGcDRC4A9St5fZzxTPNGa4_zegc4omDan4X0qg,1883
+tico/__init__.py,sha256=qQ6AhQGEdqTIol9a3pZd8krL41W7mwvr0HoA1tPCeqU,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
@@ -73,10 +73,12 @@ tico/experimental/quantization/ptq/utils/reduce_utils.py,sha256=3kWawLB91EcvvHlC
 tico/experimental/quantization/ptq/wrappers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tico/experimental/quantization/ptq/wrappers/ptq_wrapper.py,sha256=KRw_VvFJYvd2OBj4K1sYEXxUwZk9QghMw3NsgjKIAGk,1857
 tico/experimental/quantization/ptq/wrappers/quant_elementwise.py,sha256=LhEoobfvto6zKrBOKL4gmxfFFc31jHzyQV_zfps-iQM,3604
-tico/experimental/quantization/ptq/wrappers/quant_module_base.py,sha256=6RK4bn9G1pzFmkIdBdFf7liBOpb-b7rpthgD83AgkbQ,5256
-tico/experimental/quantization/ptq/wrappers/registry.py,sha256=exXl2wNNzVgC2P9gMjpF_-PqIBgYERGruzh0u1Pril0,4367
-tico/experimental/quantization/ptq/wrappers/nn/__init__.py,sha256=q4A9BiGlsa8ZdGV3y0SDiSkzkdVugsK2iz2daiJqBCY,118
+tico/experimental/quantization/ptq/wrappers/quant_module_base.py,sha256=vkcDos_knGSS29rIZuEIWkAJLHrENbGz8nCH2-iara8,5969
+tico/experimental/quantization/ptq/wrappers/registry.py,sha256=562nKSlp9qF-w4-aQeJbx2V_wMGE2FRrjIKUfRwC4Mg,4571
+tico/experimental/quantization/ptq/wrappers/nn/__init__.py,sha256=I9uTt5HfcRoMEDYHpAeATMv2TbCQiX0ZbfUFMzSJ4Qw,336
+tico/experimental/quantization/ptq/wrappers/nn/quant_layernorm.py,sha256=G5Sgt-tXnzh0Rxyk-2honmZIfEQOZlRfOsoDBdSGmA4,6887
 tico/experimental/quantization/ptq/wrappers/nn/quant_linear.py,sha256=xW-VEPB7RJoslS3xLVCdhIuMjppknvpkZleRGK4JFVQ,2240
+tico/experimental/quantization/ptq/wrappers/nn/quant_silu.py,sha256=XnJDggkWUTfXC1-BLeAbcCUtp687XLIkIIbuQlqycDw,1864
 tico/interpreter/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/interpreter/infer.py,sha256=1ZFe3DVMR2mlwBosoedqoL0-CGN_01CKLgMgxuw62KA,4861
 tico/interpreter/interpreter.py,sha256=tGbluCbrehTCqBu8mtGDNzby_ieJ2ry8_RH_eC0CQxk,3828
@@ -231,9 +233,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.dev250824.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250824.dist-info/METADATA,sha256=5PX-E9sIG566-2M_MBf8gnBXrQTSpyXzsPTeL-BLPEQ,8450
-tico-0.1.0.dev250824.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250824.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250824.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250824.dist-info/RECORD,,
+tico-0.1.0.dev250825.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250825.dist-info/METADATA,sha256=7wBkNIwJG_prscPdY7Rn_Muit4OuPN29Q8C_isHlEdI,8450
+tico-0.1.0.dev250825.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250825.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250825.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250825.dist-info/RECORD,,