tico 0.1.0.dev251106__py3-none-any.whl → 0.2.0.dev260122__py3-none-any.whl

tico/quantization/algorithm/pt2e/annotation/op/conv2d.py

@@ -1,90 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-from torch.ao.quantization.quantizer import DerivedQuantizationSpec
-
-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
-
-
-@annot_spec.register_annotator(
-    [torch.ops.aten.conv2d.default, torch.ops.aten.conv2d.padding]
-)
-def _annotate_conv2d(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    for node in gm.graph.nodes:
-        if node.op != "call_function" or node.target not in [
-            torch.ops.aten.conv2d.default,
-            torch.ops.aten.conv2d.padding,
-        ]:
-            continue
-        if filter_fn and not filter_fn(node):
-            continue
-        if quant_utils.is_annotated(node):
-            continue
-
-        args = Conv2DArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-        input_ = args.input
-        weight = args.weight
-        bias = args.bias
-
-        input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-        weight_qspec = quant_utils.get_weight_qspec(quantization_config)
-        annot_utils.annotate_input_qspec_map(node, input_, input_act_qspec)
-        annot_utils.annotate_input_qspec_map(node, weight, weight_qspec)
-        nodes_to_mark_annotated = [input_, weight, node]
-        if bias:
-
-            def _derive_bias_qparams_from_act_and_weight_qparams(obs_or_fqs):
-                act_scale, _ = obs_or_fqs[0].calculate_qparams()
-                weight_scale, _ = obs_or_fqs[1].calculate_qparams()
-                bias_scale = act_scale * weight_scale
-                bias_zero_point = torch.zeros_like(bias_scale, dtype=torch.int32)
-                return bias_scale, bias_zero_point
-
-            bias_qspec = DerivedQuantizationSpec(
-                derived_from=[
-                    (input_, node),
-                    (weight, node),
-                ],
-                derive_qparams_fn=_derive_bias_qparams_from_act_and_weight_qparams,
-                dtype=torch.int32,
-                quant_min=-(2**31),
-                quant_max=2**31 - 1,
-                qscheme=weight_qspec.qscheme,
-                ch_axis=0 if weight_qspec.qscheme == torch.per_channel_affine else None,
-            )
-            annot_utils.annotate_input_qspec_map(
-                node,
-                bias,
-                bias_qspec,
-            )
-            nodes_to_mark_annotated.append(bias)
-
-        output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-        annot_utils.annotate_output_qspec(node, output_act_qspec)
-
-        annot_utils.mark_nodes_as_annotated(nodes_to_mark_annotated)

tico/quantization/algorithm/pt2e/annotation/op/div.py

@@ -1,55 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-
-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
-
-
-@annot_spec.register_annotator([torch.ops.aten.div.Tensor])
-def _annotate_div(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    if node.op != "call_function" or node.target != torch.ops.aten.div.Tensor:
-        return
-    if filter_fn and not filter_fn(node):
-        return
-    if quant_utils.is_annotated(node):
-        return
-
-    args = DivTensorArgs(*node.args)  # type: ignore[arg-type]
-    input = args.input
-    other = args.other
-
-    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-    if isinstance(input, torch.fx.Node):
-        annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
-    if isinstance(other, torch.fx.Node):
-        annot_utils.annotate_input_qspec_map(node, other, input_act_qspec)
-
-    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-    annot_utils.annotate_output_qspec(node, output_act_qspec)
-
-    annot_utils.mark_nodes_as_annotated(node)

tico/quantization/algorithm/pt2e/annotation/op/linear.py

@@ -1,92 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-from torch.ao.quantization.quantizer import DerivedQuantizationSpec
-
-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
-
-
-@annot_spec.register_annotator([torch.ops.aten.linear.default])
-def _annotate_linear(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    if node.op != "call_function" or node.target != torch.ops.aten.linear.default:
-        return
-    if filter_fn and not filter_fn(node):
-        return
-    if quant_utils.is_annotated(node):
-        return
-
-    args = LinearArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-    input_ = args.input
-    weight = args.weight
-    bias = args.bias
-
-    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-    weight_qspec = quant_utils.get_weight_qspec(quantization_config)
-    bias_qspec = quant_utils.get_bias_qspec(quantization_config)
-
-    annot_utils.annotate_input_qspec_map(
-        node,
-        input_,
-        input_act_qspec,
-    )
-    annot_utils.annotate_input_qspec_map(
-        node,
-        weight,
-        weight_qspec,
-    )
-    nodes_to_mark_annotated = [node, weight]
-    if bias:
-
-        def _derive_bias_qparams_from_act_and_weight_qparams(obs_or_fqs):
-            act_scale, _ = obs_or_fqs[0].calculate_qparams()
-            weight_scale, _ = obs_or_fqs[1].calculate_qparams()
-            bias_scale = act_scale * weight_scale
-            bias_zero_point = torch.zeros_like(bias_scale, dtype=torch.int32)
-            return bias_scale, bias_zero_point
-
-        bias_qspec = DerivedQuantizationSpec(
-            derived_from=[
-                (input_, node),
-                (weight, node),
-            ],
-            derive_qparams_fn=_derive_bias_qparams_from_act_and_weight_qparams,
-            dtype=torch.int32,
-            quant_min=-(2**31),
-            quant_max=2**31 - 1,
-            qscheme=weight_qspec.qscheme,
-            ch_axis=0 if weight_qspec.qscheme == torch.per_channel_affine else None,
-        )
-        annot_utils.annotate_input_qspec_map(
-            node,
-            bias,
-            bias_qspec,
-        )
-        nodes_to_mark_annotated.append(bias)
-    annot_utils.annotate_output_qspec(node, output_act_qspec)
-    annot_utils.mark_nodes_as_annotated(nodes_to_mark_annotated)

tico/quantization/algorithm/pt2e/annotation/op/mean.py

@@ -1,51 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-
-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
-
-
-@annot_spec.register_annotator([torch.ops.aten.mean.dim])
-def _annotate_mean(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    if node.op != "call_function" or node.target != torch.ops.aten.mean.dim:
-        return
-    if filter_fn and not filter_fn(node):
-        return
-    if quant_utils.is_annotated(node):
-        return
-
-    args = MeanDimArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-    input = args.input
-
-    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-    annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
-
-    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-    annot_utils.annotate_output_qspec(node, output_act_qspec)
-
-    annot_utils.mark_nodes_as_annotated(node)

tico/quantization/algorithm/pt2e/annotation/op/mul.py

@@ -1,55 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-
-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
-
-
-@annot_spec.register_annotator([torch.ops.aten.mul.Tensor])
-def _annotate_mul(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    if node.op != "call_function" or node.target != torch.ops.aten.mul.Tensor:
-        return
-    if filter_fn and not filter_fn(node):
-        return
-    if quant_utils.is_annotated(node):
-        return
-
-    args = MulTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-    input = args.input
-    other = args.other
-
-    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-    if isinstance(input, torch.fx.Node):
-        annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
-    if isinstance(other, torch.fx.Node):
-        annot_utils.annotate_input_qspec_map(node, other, input_act_qspec)
-
-    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-    annot_utils.annotate_output_qspec(node, output_act_qspec)
-
-    annot_utils.mark_nodes_as_annotated(node)

tico/quantization/algorithm/pt2e/annotation/op/relu6.py

@@ -1,51 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-
-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
-
-
-@annot_spec.register_annotator([torch.ops.aten.relu6.default])
-def _annotate_relu6(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    if node.op != "call_function" or node.target != torch.ops.aten.relu6.default:
-        return
-    if filter_fn and not filter_fn(node):
-        return
-    if quant_utils.is_annotated(node):
-        return
-
-    args = Relu6Args(*node.args, **node.kwargs)  # type: ignore
-    input = args.input
-
-    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-    annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
-
-    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-    annot_utils.annotate_output_qspec(node, output_act_qspec)
-
-    annot_utils.mark_nodes_as_annotated(node)

tico/quantization/algorithm/pt2e/annotation/op/rsqrt.py

@@ -1,51 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-
-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
-
-
-@annot_spec.register_annotator([torch.ops.aten.rsqrt.default])
-def _annotate_rsqrt(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    if node.op != "call_function" or node.target != torch.ops.aten.rsqrt.default:
-        return
-    if filter_fn and not filter_fn(node):
-        return
-    if quant_utils.is_annotated(node):
-        return
-
-    args = RsqrtArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
-    input = args.input
-
-    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-    annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
-
-    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-    annot_utils.annotate_output_qspec(node, output_act_qspec)
-
-    annot_utils.mark_nodes_as_annotated(node)

tico/quantization/algorithm/pt2e/annotation/op/sub.py

@@ -1,55 +0,0 @@
-# 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 Callable, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-
-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
-
-
-@annot_spec.register_annotator([torch.ops.aten.sub.Tensor])
-def _annotate_sub(
-    gm: torch.fx.GraphModule,
-    node: torch.fx.Node,
-    quantization_config: Optional[QuantizationConfig],
-    filter_fn: Optional[Callable[[torch.fx.Node], bool]] = None,
-):
-    if node.op != "call_function" or node.target != torch.ops.aten.sub.Tensor:
-        return
-    if filter_fn and not filter_fn(node):
-        return
-    if quant_utils.is_annotated(node):
-        return
-
-    args = SubTensorArgs(*node.args)  # type: ignore[arg-type]
-    input = args.input
-    other = args.other
-
-    input_act_qspec = quant_utils.get_input_act_qspec(quantization_config)
-    if isinstance(input, torch.fx.Node):
-        annot_utils.annotate_input_qspec_map(node, input, input_act_qspec)
-    if isinstance(other, torch.fx.Node):
-        annot_utils.annotate_input_qspec_map(node, other, input_act_qspec)
-
-    output_act_qspec = quant_utils.get_output_act_qspec(quantization_config)
-    annot_utils.annotate_output_qspec(node, output_act_qspec)
-
-    annot_utils.mark_nodes_as_annotated(node)

tico/quantization/algorithm/pt2e/annotation/spec.py

@@ -1,45 +0,0 @@
-# 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 Callable, Dict, List, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-
-from tico.quantization.algorithm.pt2e.annotation.config import QuantizationConfig
-
-AnnotatorType = Callable[
-    [
-        torch.fx.GraphModule,
-        torch.fx.Node,
-        Optional[QuantizationConfig],
-        Optional[Callable[[torch.fx.Node], bool]],
-    ],
-    None,
-]
-OP_TO_ANNOTATOR: Dict[torch._ops.OpOverload, AnnotatorType] = {}
-OP_TO_SHARE_QUANT_SPEC: List[Callable] = [
-    torch.ops.aten.view_copy.default,
-    torch.ops.aten.view.default,
-]
-
-
-def register_annotator(target: List[torch._ops.OpOverload]):
-    def decorator(annotator: AnnotatorType):
-        for t in target:
-            OP_TO_ANNOTATOR[t] = annotator
-        return annotator
-
-    return decorator

tico/quantization/algorithm/pt2e/annotation/utils.py

@@ -1,88 +0,0 @@
-# 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 List, Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
-from torch.ao.quantization.quantizer import (
-    QuantizationAnnotation,
-    SharedQuantizationSpec,
-)
-
-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):
-    quantization_annotation: QuantizationAnnotation = node.meta.get(
-        "quantization_annotation", QuantizationAnnotation()
-    )
-    quantization_annotation.input_qspec_map[input_node] = qspec
-    node.meta["quantization_annotation"] = quantization_annotation
-
-
-def annotate_output_qspec(node: torch.fx.Node, qspec):
-    quantization_annotation: QuantizationAnnotation = node.meta.get(
-        "quantization_annotation", QuantizationAnnotation()
-    )
-    quantization_annotation.output_qspec = qspec
-    node.meta["quantization_annotation"] = quantization_annotation
-
-
-def mark_nodes_as_annotated(nodes: List[torch.fx.Node] | torch.fx.Node):
-    if isinstance(nodes, torch.fx.Node):
-        nodes = [nodes]
-    for node in nodes:
-        if node is not None:
-            if "quantization_annotation" not in node.meta:
-                node.meta["quantization_annotation"] = QuantizationAnnotation()
-            node.meta["quantization_annotation"]._annotated = True
-
-
-def propagate_annotation_forward(model: torch.fx.GraphModule) -> None:
-    for n in model.graph.nodes:
-        if n.op != "call_function" or n.target not in annot_spec.OP_TO_SHARE_QUANT_SPEC:
-            continue
-
-        prev_node = n.args[0]
-        if not isinstance(prev_node, torch.fx.Node):
-            continue
-
-        quantization_annotation: Optional[QuantizationAnnotation] = prev_node.meta.get(
-            "quantization_annotation", None
-        )
-        if not quantization_annotation:
-            continue
-
-        output_qspec = quantization_annotation.output_qspec
-        if not output_qspec:
-            continue
-
-        # Make sure current node is not annotated
-        if (
-            "quantization_annotation" in n.meta
-            and n.meta["quantization_annotation"]._annotated
-        ):
-            continue
-
-        shared_qspec = SharedQuantizationSpec(prev_node)
-        # Propagate the previous output_qspec to the current node
-        n.meta["quantization_annotation"] = QuantizationAnnotation(
-            input_qspec_map={
-                prev_node: shared_qspec,
-            },
-            output_qspec=shared_qspec,
-            _annotated=True,
-        )