Ease Core ML partitioner and quantizer imports.

backends/apple/coreml/README.md

@@ -28,7 +28,7 @@ import torch
 import executorch.exir
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
-from executorch.backends.apple.coreml.partition.coreml_partitioner import CoreMLPartitioner
+from executorch.backends.apple.coreml.partition import CoreMLPartitioner
 
 class Model(torch.nn.Module):
     def __init__(self):
@@ -72,7 +72,7 @@ from torch.ao.quantization.quantize_pt2e import (
     prepare_qat_pt2e,
 )
 
-from executorch.backends.apple.coreml.quantizer.coreml_quantizer import CoreMLQuantizer
+from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer
 from coremltools.optimize.torch.quantization.quantization_config import (
     LinearQuantizerConfig,
     QuantizationScheme,

backends/apple/coreml/partition/__init__.py

@@ -0,0 +1,9 @@
+# Copyright © 2023 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+from .coreml_partitioner import CoreMLPartitioner
+
+__all__ = [
+    CoreMLPartitioner,
+]

backends/apple/coreml/quantizer/__init__.py

@@ -0,0 +1,9 @@
+# Copyright © 2023 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+from .coreml_quantizer import CoreMLQuantizer
+
+__all__ = [
+    CoreMLQuantizer,
+]

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -9,9 +9,7 @@
 import torch
 import torchvision
 
-from executorch.backends.apple.coreml.partition.coreml_partitioner import (
-    CoreMLPartitioner,
-)
+from executorch.backends.apple.coreml.partition import CoreMLPartitioner
 
 
 class TestCoreMLPartitioner(unittest.TestCase):

backends/apple/coreml/test/test_coreml_quantizer.py

@@ -14,7 +14,7 @@
     QuantizationScheme,
 )
 
-from executorch.backends.apple.coreml.quantizer.coreml_quantizer import CoreMLQuantizer
+from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer
 from torch._export import capture_pre_autograd_graph
 from torch.ao.quantization.quantize_pt2e import (
     convert_pt2e,

examples/apple/coreml/scripts/export.py

@@ -16,9 +16,7 @@
 
 from executorch.backends.apple.coreml.compiler import CoreMLBackend
 
-from executorch.backends.apple.coreml.partition.coreml_partitioner import (
-    CoreMLPartitioner,
-)
+from executorch.backends.apple.coreml.partition import CoreMLPartitioner
 from executorch.exir import to_edge
 
 from executorch.exir.backend.backend_api import to_backend

examples/models/llama2/lib/partitioner_lib.py

@@ -57,16 +57,14 @@ def get_coreml_partitioner(args):
         args.use_kv_cache is True
     ), "CoreML backend currently only supports static shape and use_kv_cache=True is the only way to support it at the moment"
     try:
-        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.partition.coreml_partitioner`.
+        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `coremltools`.
         import coremltools as ct
 
         # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.compiler`
         from executorch.backends.apple.coreml.compiler import CoreMLBackend
 
-        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.partition.coreml_partitioner`
-        from executorch.backends.apple.coreml.partition.coreml_partitioner import (
-            CoreMLPartitioner,
-        )
+        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.partition`
+        from executorch.backends.apple.coreml.partition import CoreMLPartitioner
     except ImportError:
         raise ImportError(
             "Please install the CoreML backend follwing https://pytorch.org/executorch/main/build-run-coreml.html"