up

Author: metascroy

backends/arm/quantizer/arm_quantizer.py

@@ -31,23 +31,22 @@
 )  # usort: skip
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from torch.fx import GraphModule, Node
-from torchao.quantization.pt2e import _ObserverOrFakeQuantizeConstructor
-from torchao.quantization.pt2e.fake_quantize import (
+from torchao.quantization.pt2e import (
     FakeQuantize,
     FusedMovingAvgObsFakeQuantize,
-)
-from torchao.quantization.pt2e.observer import (
     HistogramObserver,
     MinMaxObserver,
     MovingAverageMinMaxObserver,
     MovingAveragePerChannelMinMaxObserver,
+    ObserverOrFakeQuantizeConstructor,
     PerChannelMinMaxObserver,
     PlaceholderObserver,
 )
-from torchao.quantization.pt2e.quantizer import QuantizationSpec, Quantizer
-from torchao.quantization.pt2e.quantizer.utils import (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
+from torchao.quantization.pt2e.quantizer import (
+    annotate_input_qspec_map,
+    annotate_output_qspec,
+    QuantizationSpec,
+    Quantizer,
 )
 
 __all__ = [
@@ -97,7 +96,7 @@ def get_symmetric_quantization_config(
     weight_qscheme = (
         torch.per_channel_symmetric if is_per_channel else torch.per_tensor_symmetric
     )
-    weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = (
+    weight_observer_or_fake_quant_ctr: ObserverOrFakeQuantizeConstructor = (
         MinMaxObserver
     )
     if is_qat:
@@ -337,14 +336,14 @@ def _annotate_io(
             if is_annotated(node):
                 continue
             if node.op == "placeholder" and len(node.users) > 0:
-                _annotate_output_qspec(
+                annotate_output_qspec(
                     node,
                     quantization_config.get_output_act_qspec(),
                 )
                 mark_node_as_annotated(node)
             if node.op == "output":
                 parent = node.all_input_nodes[0]
-                _annotate_input_qspec_map(
+                annotate_input_qspec_map(
                     node, parent, quantization_config.get_input_act_qspec()
                 )
                 mark_node_as_annotated(node)

backends/arm/quantizer/quantization_annotator.py

@@ -14,13 +14,11 @@
 from executorch.backends.arm.tosa_utils import get_node_debug_info
 from torch.fx import Node
 from torchao.quantization.pt2e.quantizer import (
+    annotate_input_qspec_map,
+    annotate_output_qspec,
     QuantizationSpecBase,
     SharedQuantizationSpec,
 )
-from torchao.quantization.pt2e.quantizer.utils import (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
-)
 
 from .arm_quantizer_utils import (
     is_annotated,
@@ -121,7 +119,7 @@ def _annotate_input(node: Node, quant_property: _QuantProperty):
         strict=True,
     ):
         assert isinstance(n_arg, Node)
-        _annotate_input_qspec_map(node, n_arg, qspec)
+        annotate_input_qspec_map(node, n_arg, qspec)
         if quant_property.mark_annotated:
             mark_node_as_annotated(n_arg)  # type: ignore[attr-defined]
 
@@ -132,7 +130,7 @@ def _annotate_output(node: Node, quant_property: _QuantProperty):
     assert not quant_property.optional
     assert quant_property.index == 0, "Only one output annotation supported currently"
 
-    _annotate_output_qspec(node, quant_property.qspec)
+    annotate_output_qspec(node, quant_property.qspec)
 
 
 def _match_pattern(

backends/example/example_backend_delegate_passes/permute_memory_formats_pass.py

@@ -11,7 +11,7 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dim_order_utils import get_dim_order
 from executorch.exir.pass_base import ExportPass, PassResult
-from torchao.quantization.pt2e.pt2e.graph_utils import find_sequential_partitions
+from torchao.quantization.pt2e import find_sequential_partitions
 
 
 class PermuteMemoryFormatsPass(ExportPass):

backends/example/example_partitioner.py

@@ -21,7 +21,7 @@
 from executorch.exir.graph_module import get_control_flow_submodules
 from torch.export import ExportedProgram
 from torch.fx.passes.operator_support import OperatorSupportBase
-from torchao.quantization.pt2e.pt2e.graph_utils import find_sequential_partitions
+from torchao.quantization.pt2e import find_sequential_partitions
 
 
 @final

backends/example/example_quantizer.py

@@ -11,8 +11,8 @@
 from executorch.backends.example.example_operators.ops import module_to_annotator
 from executorch.backends.xnnpack.quantizer.xnnpack_quantizer_utils import OperatorConfig
 from torch import fx
+from torchao.quantization.pt2e import find_sequential_partitions
 from torchao.quantization.pt2e.observer import HistogramObserver, MinMaxObserver
-from torchao.quantization.pt2e.pt2e.graph_utils import find_sequential_partitions
 from torchao.quantization.pt2e.quantizer import QuantizationSpec, Quantizer
 
 

backends/mediatek/quantizer/annotator.py

@@ -16,10 +16,10 @@
     SubgraphMatcherWithNameNodeMap,
 )
 
-from torchao.quantization.pt2e.quantizer import QuantizationAnnotation
-from torchao.quantization.pt2e.quantizer.utils import (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
+from torchao.quantization.pt2e.quantizer import (
+    annotate_input_qspec_map,
+    annotate_output_qspec as _annotate_output_qspec,
+    QuantizationAnnotation,
 )
 
 from .qconfig import QuantizationConfig
@@ -108,7 +108,7 @@ def _annotate_fused_activation_pattern(
             torch.ops.aten.linear.default,
         ]:
             weight_node = producer_node.args[1]
-            _annotate_input_qspec_map(
+            annotate_input_qspec_map(
                 producer_node,
                 weight_node,
                 quant_config.weight,
@@ -201,7 +201,7 @@ def annotate_affine_ops(node: Node, quant_config: QuantizationConfig) -> None:
         return
 
     weight_node = node.args[1]
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quant_config.weight,
@@ -260,5 +260,5 @@ def annotate_embedding_op(node: Node, quant_config: QuantizationConfig) -> None:
         return
 
     wgt_node = node.args[0]
-    _annotate_input_qspec_map(node, wgt_node, quant_config.activation)
+    annotate_input_qspec_map(node, wgt_node, quant_config.activation)
     _mark_as_annotated([node])

backends/qualcomm/quantizer/annotators.py

@@ -23,8 +23,8 @@
     SharedQuantizationSpec,
 )
 from torchao.quantization.pt2e.quantizer.utils import (
-    _annotate_input_qspec_map,
-    _annotate_output_qspec,
+    annotate_input_qspec_map,
+    annotate_output_qspec,
 )
 
 from .qconfig import (
@@ -618,19 +618,19 @@ def annotate_rms_norm(node: Node, quantization_config: QuantizationConfig) -> No
         return
 
     # TODO current only support 16a16w
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         quantization_config.input_activation,
     )
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quantization_config.input_activation,
     )
     nodes_to_mark_annotated = [node]
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
 
@@ -819,25 +819,25 @@ def annotate_group_norm(node: Node, quantization_config: QuantizationConfig) ->
     if _is_annotated([node]):
         return
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         quantization_config.input_activation,
     )
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quantization_config.weight,
     )
     nodes_to_mark_annotated = [node, weight_node]
     if bias_node:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             bias_node,
             quantization_config.bias,
         )
         nodes_to_mark_annotated.append(bias_node)
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
 
@@ -1002,12 +1002,12 @@ def annotate_linear(node: Node, quantization_config: QuantizationConfig) -> None
     if _is_annotated([node]):
         return
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         quantization_config.input_activation,
     )
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         weight_node,
         quantization_config.weight,
@@ -1018,9 +1018,9 @@ def annotate_linear(node: Node, quantization_config: QuantizationConfig) -> None
             bias_config = quantization_config.bias(node)
         else:
             bias_config = quantization_config.bias
-        _annotate_input_qspec_map(node, bias_node, bias_config)
+        annotate_input_qspec_map(node, bias_node, bias_config)
         nodes_to_mark_annotated.append(bias_node)
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
     # We use get_source_partition in pass, but it is the same source for MultiheadAttention, so we need to change its source_fn_stack.
@@ -1038,29 +1038,29 @@ def annotate_batch_and_instance_norm(
         return
 
     annotated_args = [act]
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act,
         quantization_config.input_activation,
     )
     # QNN requires uint8 instead of int8 in 'weight' config
     if weight is not None:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             weight,
             quantization_config.input_activation,
         )
         annotated_args.append(weight)
 
     if bias is not None:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             bias,
             quantization_config.bias,
         )
         annotated_args.append(bias)
 
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated([node, *annotated_args])
 
 
@@ -1070,7 +1070,7 @@ def annotate_getitem(node: Node, quantization_config: QuantizationConfig) -> Non
         return
 
     if _is_float_tensor(node):
-        _annotate_output_qspec(node, quantization_config.output_activation)
+        annotate_output_qspec(node, quantization_config.output_activation)
         _mark_nodes_as_annotated([node])
 
 
@@ -1086,32 +1086,32 @@ def annotate_layer_norm(node: Node, quantization_config: QuantizationConfig) ->
         return
     input_act_qspec = quantization_config.input_activation
 
-    _annotate_input_qspec_map(
+    annotate_input_qspec_map(
         node,
         act_node,
         input_act_qspec,
     )
     if input_act_qspec.dtype == torch.int32:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             weight_node,
             get_16a16w_qnn_ptq_config().weight,
         )
     else:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             weight_node,
             input_act_qspec,
         )
     nodes_to_mark_annotated = [node, weight_node]
     if bias_node:
-        _annotate_input_qspec_map(
+        annotate_input_qspec_map(
             node,
             bias_node,
             quantization_config.bias,
         )
         nodes_to_mark_annotated.append(bias_node)
-    _annotate_output_qspec(node, quantization_config.output_activation)
+    annotate_output_qspec(node, quantization_config.output_activation)
     _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
 

backends/qualcomm/quantizer/observers/per_block_param_observer.py

@@ -7,12 +7,12 @@
 from typing import Tuple
 
 import torch
-from torchao.quantization.pt2e.observer import MappingType, PerBlock
-from torchao.quantization.pt2e.pt2e._affine_quantization import (
+from torchao.quantization.pt2e._affine_quantization import (
     _get_reduction_params,
     AffineQuantizedMinMaxObserver,
     choose_qparams_affine_with_min_max,
 )
+from torchao.quantization.pt2e.observer import MappingType, PerBlock
 
 
 class PerBlockParamObserver(AffineQuantizedMinMaxObserver):

backends/transforms/duplicate_dynamic_quant_chain.py

@@ -12,9 +12,9 @@
 from torch.fx.node import map_arg
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 
-from torchao.quantization.pt2e.pt2e.utils import (
+from torchao.quantization.pt2e.utils import (
     _filter_sym_size_users,
-    _is_valid_annotation,
+    _is_valid_annotation,  # @nocommit not found
 )
 
 

backends/xnnpack/quantizer/xnnpack_quantizer.py

@@ -30,12 +30,12 @@
     PlaceholderObserver,
 )
 from torchao.quantization.pt2e.quantizer import QuantizationSpec, Quantizer
-from torchao.quantization.pt2e.quantizer.utils import _get_module_name_filter
+from torchao.quantization.pt2e.quantizer.utils import get_module_name_filter
 
 
 if TYPE_CHECKING:
     from torch.fx import Node
-    from torchao.quantization.pt2e import _ObserverOrFakeQuantizeConstructor
+    from torchao.quantization.pt2e import ObserverOrFakeQuantizeConstructor
 
 
 __all__ = [
@@ -140,7 +140,7 @@ def get_symmetric_quantization_config(
     weight_qscheme = (
         torch.per_channel_symmetric if is_per_channel else torch.per_tensor_symmetric
     )
-    weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = (
+    weight_observer_or_fake_quant_ctr: ObserverOrFakeQuantizeConstructor = (
         MinMaxObserver
     )
     if is_qat:
@@ -228,7 +228,7 @@ def _get_not_module_type_or_name_filter(
     tp_list: list[Callable], module_name_list: list[str]
 ) -> Callable[[Node], bool]:
     module_type_filters = [_get_module_type_filter(tp) for tp in tp_list]
-    module_name_list_filters = [_get_module_name_filter(m) for m in module_name_list]
+    module_name_list_filters = [get_module_name_filter(m) for m in module_name_list]
 
     def not_module_type_or_name_filter(n: Node) -> bool:
         return not any(f(n) for f in module_type_filters + module_name_list_filters)
@@ -421,7 +421,7 @@ def _annotate_for_quantization_config(
         module_name_list = list(self.module_name_config.keys())
         for module_name, config in self.module_name_config.items():
             self._annotate_all_patterns(
-                model, config, _get_module_name_filter(module_name)
+                model, config, get_module_name_filter(module_name)
             )
 
         tp_list = list(self.module_type_config.keys())