[1/n] Move a few utils under pt2e/pt2e/ to pt2e/ (#2083)

* [1/n] Move pt2e/pt2e/utils.py (and graph, export, qat utils) to pt2e/

Summary:
#2082

Test Plan:
pytest test/quantization/pt2e

Reviewers:

Subscribers:

Tasks:

Tags:

* fix

Author: jerryzh168

test/quantization/pt2e/test_graph_utils.py

@@ -12,7 +12,7 @@
 import torch._dynamo as torchdynamo
 from torch.testing._internal.common_utils import IS_WINDOWS, TestCase, run_tests
 
-from torchao.quantization.pt2e.pt2e.graph_utils import (
+from torchao.quantization.pt2e.graph_utils import (
     find_sequential_partitions,
     get_equivalent_types,
     update_equivalent_types_dict,

test/quantization/pt2e/test_numeric_debugger.py

@@ -22,7 +22,7 @@
     generate_numeric_debug_handle,
     prepare_for_propagation_comparison,
 )
-from torchao.quantization.pt2e.pt2e.graph_utils import bfs_trace_with_node_process
+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 (
     XNNPACKQuantizer,

test/quantization/pt2e/test_quantize_pt2e.py

@@ -2077,14 +2077,12 @@ def test_model_is_exported(self):
         exported_gm = export_for_training(m, example_inputs, strict=True).module()
         fx_traced_gm = torch.fx.symbolic_trace(m, example_inputs)
         self.assertTrue(
-            torchao.quantization.pt2e.pt2e.export_utils.model_is_exported(exported_gm)
+            torchao.quantization.pt2e.export_utils.model_is_exported(exported_gm)
         )
         self.assertFalse(
-            torchao.quantization.pt2e.pt2e.export_utils.model_is_exported(fx_traced_gm)
-        )
-        self.assertFalse(
-            torchao.quantization.pt2e.pt2e.export_utils.model_is_exported(m)
+            torchao.quantization.pt2e.export_utils.model_is_exported(fx_traced_gm)
         )
+        self.assertFalse(torchao.quantization.pt2e.export_utils.model_is_exported(m))
 
     def test_reentrant(self):
         """Test we can safely call quantization apis multiple times"""

torchao/quantization/pt2e/__init__.py

@@ -5,6 +5,15 @@
 import torch
 from torch import Tensor
 
+from torchao.quantization.pt2e.export_utils import (
+    _allow_exported_model_train_eval as allow_exported_model_train_eval,
+)
+from torchao.quantization.pt2e.export_utils import (
+    _move_exported_model_to_eval as move_exported_model_to_eval,
+)
+from torchao.quantization.pt2e.export_utils import (
+    _move_exported_model_to_train as move_exported_model_to_train,
+)
 from torchao.quantization.pt2e.pt2e._numeric_debugger import (  # noqa: F401
     CUSTOM_KEY,
     NUMERIC_DEBUG_HANDLE_KEY,
@@ -13,15 +22,6 @@
     generate_numeric_debug_handle,
     prepare_for_propagation_comparison,
 )
-from torchao.quantization.pt2e.pt2e.export_utils import (
-    _allow_exported_model_train_eval as allow_exported_model_train_eval,
-)
-from torchao.quantization.pt2e.pt2e.export_utils import (
-    _move_exported_model_to_eval as move_exported_model_to_eval,
-)
-from torchao.quantization.pt2e.pt2e.export_utils import (
-    _move_exported_model_to_train as move_exported_model_to_train,
-)
 
 from .fake_quantize import (
     FakeQuantize,

torchao/quantization/pt2e/pt2e/export_utils.py renamed to torchao/quantization/pt2e/export_utils.py

@@ -16,7 +16,7 @@
 from torch.fx import GraphModule, Node
 from torch.nn import functional as F
 
-from torchao.quantization.pt2e.pt2e.graph_utils import bfs_trace_with_node_process
+from torchao.quantization.pt2e.graph_utils import bfs_trace_with_node_process
 
 NUMERIC_DEBUG_HANDLE_KEY = "numeric_debug_handle"
 CUSTOM_KEY = "custom"

torchao/quantization/pt2e/pt2e/graph_utils.py renamed to torchao/quantization/pt2e/graph_utils.py

@@ -12,10 +12,12 @@
 from torch.fx.node import map_arg
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 
-from torchao.quantization.pt2e.pt2e.utils import (
-    _filter_sym_size_users,
+from torchao.quantization.pt2e.quantizer.utils import (
     _is_valid_annotation,
 )
+from torchao.quantization.pt2e.utils import (
+    _filter_sym_size_users,
+)
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)

torchao/quantization/pt2e/pt2e/_numeric_debugger.py

@@ -12,12 +12,14 @@
 from torch._export.error import InternalError
 from torch.fx.passes.infra.pass_base import PassBase, PassResult
 
-from torchao.quantization.pt2e.pt2e.utils import (
+from torchao.quantization.pt2e.quantizer import QuantizationSpecBase
+from torchao.quantization.pt2e.quantizer.utils import (
+    _is_valid_annotation,
+)
+from torchao.quantization.pt2e.utils import (
     _filter_sym_size_users,
     _find_q_dq_node_for_user,
-    _is_valid_annotation,
 )
-from torchao.quantization.pt2e.quantizer import QuantizationSpecBase
 from torchao.quantization.quant_primitives import quant_lib  # noqa: F401
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 

torchao/quantization/pt2e/pt2e/duplicate_dq_pass.py

@@ -16,8 +16,8 @@
 from torch.fx import GraphModule
 from torch.fx.subgraph_rewriter import replace_pattern
 
-from torchao.quantization.pt2e.pt2e.export_utils import _WrapperModule
-from torchao.quantization.pt2e.pt2e.utils import (
+from torchao.quantization.pt2e.export_utils import _WrapperModule
+from torchao.quantization.pt2e.utils import (
     _get_aten_graph_module_for_pattern,
     _replace_literals_with_existing_placeholders,
     _replace_literals_with_new_placeholders,