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

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/adaptive_avg_pool2d.py

@@ -19,12 +19,10 @@ if TYPE_CHECKING:
 import torch
 from torch.ao.quantization.quantizer import SharedQuantizationSpec
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import AdaptiveAvgPool2dArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/add.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import AddTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/conv2d.py

@@ -19,12 +19,10 @@ if TYPE_CHECKING:
 import torch
 from torch.ao.quantization.quantizer import DerivedQuantizationSpec
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import Conv2DArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/div.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import DivTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/linear.py

@@ -19,12 +19,10 @@ if TYPE_CHECKING:
 import torch
 from torch.ao.quantization.quantizer import DerivedQuantizationSpec
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import LinearArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/mean.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import MeanDimArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/mul.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import MulTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/relu6.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import Relu6Args
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/rsqrt.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import RsqrtArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/op/sub.py

@@ -18,12 +18,10 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
-import tico.experimental.quantization.algorithm.pt2e.annotation.utils as annot_utils
-import tico.experimental.quantization.algorithm.pt2e.utils as quant_utils
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.utils as annot_utils
+import tico.quantization.algorithm.pt2e.utils as quant_utils
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 from tico.utils.validate_args_kwargs import SubTensorArgs
 
 

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/spec.py

@@ -18,9 +18,7 @@ if TYPE_CHECKING:
     import torch.fx
 import torch
 
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 
 AnnotatorType = Callable[
     [

tico/{experimental/quantization → quantization}/algorithm/pt2e/annotation/utils.py

@@ -22,7 +22,7 @@ from torch.ao.quantization.quantizer import (
     SharedQuantizationSpec,
 )
 
-import tico.experimental.quantization.algorithm.pt2e.annotation.spec as annot_spec
+import tico.quantization.algorithm.pt2e.annotation.spec as annot_spec
 
 
 def annotate_input_qspec_map(node: torch.fx.Node, input_node: torch.fx.Node, qspec):

tico/{experimental/quantization → quantization}/algorithm/pt2e/quantizer.py

@@ -18,13 +18,16 @@ import torch
 
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 
-from tico.experimental.quantization.algorithm.pt2e.annotation.annotator import (
+from tico.quantization.algorithm.pt2e.annotation.annotator import (
     get_asymmetric_quantization_config,
     PT2EAnnotator,
 )
-from tico.experimental.quantization.quantizer import BaseQuantizer
+from tico.quantization.config.pt2e import PT2EConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
 
 
+@register_quantizer(PT2EConfig)
 class PT2EQuantizer(BaseQuantizer):
     """
     Quantizer for applying pytorch 2.0 export quantization (typically for activation quantization).

tico/{experimental/quantization → quantization}/algorithm/pt2e/utils.py

@@ -20,9 +20,7 @@ import torch
 from torch.ao.quantization.quantizer import QuantizationSpec
 from torch.ao.quantization.quantizer.utils import _get_module_name_filter
 
-from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
-    QuantizationConfig,
-)
+from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
 
 
 def get_module_type_filter(tp: Callable):

tico/{experimental/quantization → quantization}/algorithm/smoothquant/observer.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Literal
 
 import torch
 
@@ -21,18 +21,24 @@ import torch
 class ChannelwiseMaxActsObserver:
     """
     Observer to calcuate channelwise maximum activation
+    It supports collecting activations from either module inputs or outputs.
     """
 
-    def __init__(self, model):
+    def __init__(
+        self, model: torch.nn.Module, acts_from: Literal["input", "output"] = "input"
+    ):
         """
         model
             A torch module whose activations are to be analyzed.
+        acts_from
+            Where to hook: "input" for forward-pre-hook, "output" for forward-hook.
         hooks
-            A list to store the hooks which are registered to collect activation statistics.
+            A list to store the hooks registered to collect activation statistics.
         max_acts
-            A dictionary to store the maximum activation values
+            A dictionary to store the per-channel maxima.
         """
         self.model = model
+        self.acts_from: Literal["input", "output"] = acts_from
         self.hooks: List[Any] = []
         self.max_acts: Dict[str, torch.Tensor] = {}
 
@@ -62,13 +68,25 @@ class ChannelwiseMaxActsObserver:
                 input = input[0]
             stat_tensor(name, input)
 
+        def stat_output_hook(m, input, output, name):
+            if isinstance(output, tuple):
+                output = output[0]
+            stat_tensor(name, output)
+
         for name, m in self.model.named_modules():
             if isinstance(m, torch.nn.Linear):
-                self.hooks.append(
-                    m.register_forward_pre_hook(
-                        functools.partial(stat_input_hook, name=name)
+                if self.acts_from == "input":
+                    self.hooks.append(
+                        m.register_forward_pre_hook(
+                            functools.partial(stat_input_hook, name=name)
+                        )
+                    )
+                else:  # "output"
+                    self.hooks.append(
+                        m.register_forward_hook(
+                            functools.partial(stat_output_hook, name=name)
+                        )
                     )
-                )
 
     def remove(self):
         for hook in self.hooks:

tico/{experimental/quantization → quantization}/algorithm/smoothquant/quantizer.py

@@ -16,20 +16,37 @@ from typing import Any, Dict, Optional
 
 import torch
 
-from tico.experimental.quantization.algorithm.smoothquant.observer import (
-    ChannelwiseMaxActsObserver,
-)
+from tico.quantization.algorithm.smoothquant.observer import ChannelwiseMaxActsObserver
 
-from tico.experimental.quantization.algorithm.smoothquant.smooth_quant import (
-    apply_smoothing,
-)
-from tico.experimental.quantization.config import SmoothQuantConfig
-from tico.experimental.quantization.quantizer import BaseQuantizer
+from tico.quantization.algorithm.smoothquant.smooth_quant import apply_smoothing
+from tico.quantization.config.smoothquant import SmoothQuantConfig
+from tico.quantization.quantizer import BaseQuantizer
+from tico.quantization.quantizer_registry import register_quantizer
 
 
+@register_quantizer(SmoothQuantConfig)
 class SmoothQuantQuantizer(BaseQuantizer):
     """
     Quantizer for applying the SmoothQuant algorithm
+
+    Q) Why allow choosing between input and output activations?
+
+    SmoothQuant relies on channel-wise activation statistics to balance
+    weights and activations. In practice, there are two natural sources:
+
+    - "input": captures the tensor right before a Linear layer
+      (forward-pre-hook). This matches the original SmoothQuant paper
+      and focuses on scaling the raw hidden state.
+
+    - "output": captures the tensor right after a Linear layer
+      (forward-hook). This can better reflect post-weight dynamics,
+      especially when subsequent operations (bias, activation functions)
+      dominate the dynamic range.
+
+    Allowing both options provides flexibility: depending on model
+    architecture and calibration data, one may yield lower error than
+    the other. The default remains "input" for compatibility, but "output"
+    can be selected to empirically reduce error or runtime overhead.
     """
 
     def __init__(self, config: SmoothQuantConfig):
@@ -37,6 +54,7 @@ class SmoothQuantQuantizer(BaseQuantizer):
 
         self.alpha = config.alpha
         self.custom_alpha_map = config.custom_alpha_map
+        self.acts_from = config.acts_from  # "input" (default) or "output"
         self.observer: Optional[ChannelwiseMaxActsObserver] = None
 
     @torch.no_grad()
@@ -55,7 +73,8 @@ class SmoothQuantQuantizer(BaseQuantizer):
         Returns:
             The model prepared for SmoothQuant quantization.
         """
-        self.observer = ChannelwiseMaxActsObserver(model)
+        # Attach hooks according to `config.acts_from`
+        self.observer = ChannelwiseMaxActsObserver(model, acts_from=self.acts_from)
         self.observer.attach()
 
         return model