[3/n] Flatten import path (#2125)

* [3/n] Flatten import path

Summary:
We only keep two import paths:
```
torchao.quantization.pt2e
torchao.quantization.pt2e.quantizer
```

for all classes and utils.

Next: remove the underscore before public util functions

Test Plan:
pytest test/quantization/pt2e

Reviewers:

Subscribers:

Tasks:

Tags:

* move refactor

* fix circular import

* import

Author: jerryzh168

test/quantization/pt2e/test_duplicate_dq.py

@@ -26,10 +26,10 @@
     Quantizer,
     SharedQuantizationSpec,
 )
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer import (
+from torchao.testing.pt2e._xnnpack_quantizer import (
     get_symmetric_quantization_config,
 )
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer_utils import (
+from torchao.testing.pt2e._xnnpack_quantizer_utils import (
     OP_TO_ANNOTATOR,
     QuantizationConfig,
 )

test/quantization/pt2e/test_metadata_porting.py

@@ -16,10 +16,10 @@
 
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 from torchao.quantization.pt2e.quantizer import QuantizationAnnotation, Quantizer
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer import (
+from torchao.testing.pt2e._xnnpack_quantizer import (
     get_symmetric_quantization_config,
 )
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer_utils import OP_TO_ANNOTATOR
+from torchao.testing.pt2e._xnnpack_quantizer_utils import OP_TO_ANNOTATOR
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_7
 
 

test/quantization/pt2e/test_numeric_debugger.py

@@ -24,7 +24,7 @@
 )
 from torchao.quantization.pt2e.graph_utils import bfs_trace_with_node_process
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer import (
+from torchao.testing.pt2e._xnnpack_quantizer import (
     XNNPACKQuantizer,
     get_symmetric_quantization_config,
 )

test/quantization/pt2e/test_quantize_pt2e.py

@@ -57,11 +57,11 @@
 from torchao.quantization.pt2e.quantizer.embedding_quantizer import (  # noqa: F811
     EmbeddingQuantizer,
 )
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer import (
+from torchao.testing.pt2e._xnnpack_quantizer import (
     XNNPACKQuantizer,
     get_symmetric_quantization_config,
 )
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer_utils import (
+from torchao.testing.pt2e._xnnpack_quantizer_utils import (
     OP_TO_ANNOTATOR,
     QuantizationConfig,
 )
@@ -1328,6 +1328,40 @@ def validate(self, model: torch.fx.GraphModule) -> None:
         with self.assertRaises(Exception):
             m = prepare_pt2e(m, BackendAQuantizer())
 
+    def _quantize(self, m, quantizer, example_inputs, is_qat: bool = False):
+        # resetting dynamo cache
+        torch._dynamo.reset()
+
+        m = export_for_training(
+            m,
+            example_inputs,
+        ).module()
+        if is_qat:
+            m = prepare_qat_pt2e(m, quantizer)
+        else:
+            m = prepare_pt2e(m, quantizer)
+        m(*example_inputs)
+        m = convert_pt2e(m)
+        return m
+
+    def _get_pt2e_quantized_linear(self, is_per_channel=False) -> torch.fx.GraphModule:
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.linear = torch.nn.Linear(2, 2)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        quantizer = XNNPACKQuantizer()
+        operator_config = get_symmetric_quantization_config(
+            is_per_channel=is_per_channel
+        )
+        quantizer.set_global(operator_config)
+        example_inputs = (torch.randn(2, 2),)
+        m = M().eval()
+        return self._quantize(m, quantizer, example_inputs)
+
     def test_fold_quantize(self):
         """Test to make sure the quantized model gets quantized weight (quantize_per_tensor op is folded)"""
         m = self._get_pt2e_quantized_linear()

test/quantization/pt2e/test_quantize_pt2e_qat.py

@@ -47,7 +47,7 @@
     QuantizationSpec,
     Quantizer,
 )
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer import (
+from torchao.testing.pt2e._xnnpack_quantizer import (
     XNNPACKQuantizer,
     get_symmetric_quantization_config,
 )

test/quantization/pt2e/test_representation.py

@@ -23,7 +23,7 @@
 
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 from torchao.quantization.pt2e.quantizer import Quantizer
-from torchao.quantization.pt2e.quantizer.xnnpack_quantizer import (
+from torchao.testing.pt2e._xnnpack_quantizer import (
     XNNPACKQuantizer,
     get_symmetric_quantization_config,
 )