Rename to aot_types.py

LiteRT-PiperOrigin-RevId: 883235059

Author: gcarranza-1

ci/tools/python/wheel/BUILD

@@ -115,7 +115,7 @@ ALL_PY_SRC_MODULES = [
     "//litert/python/aot:aot_compile",
     "//litert/python/aot/ai_pack:export_lib",
     "//litert/python/aot/core:common",
-    "//litert/python/aot/core:types",
+    "//litert/python/aot/core:aot_types",
     "//litert/python/aot/core:apply_plugin",
     "//litert/python/aot/core:mlir_transforms",
     "//litert/python/aot/core:components",

litert/python/aot/BUILD

@@ -25,10 +25,10 @@ pytype_strict_library(
     visibility = ["//visibility:public"],
     deps = [
         ":prepare_for_npu",
+        "//litert/python/aot/core:aot_types",
         "//litert/python/aot/core:apply_plugin",
         "//litert/python/aot/core:components",
         "//litert/python/aot/core:mlir_transforms",
-        "//litert/python/aot/core:types",
         "//litert/python/aot/vendors:import_vendor",
     ],
 )
@@ -51,9 +51,9 @@ pytype_strict_contrib_test(
     ],
     deps = [
         ":aot_compile",
+        "//litert/python/aot/core:aot_types",
         "//litert/python/aot/core:common",
         "//litert/python/aot/core:test_common",
-        "//litert/python/aot/core:types",
         "//litert/python/aot/vendors/mediatek:target",
         "//litert/python/aot/vendors/qualcomm:target",
         "@absl_py//absl/testing:absltest",
@@ -64,9 +64,9 @@ pytype_strict_library(
     name = "prepare_for_npu",
     srcs = ["prepare_for_npu.py"],
     deps = [
+        "//litert/python/aot/core:aot_types",
         "//litert/python/aot/core:common",
         "//litert/python/aot/core:components",
-        "//litert/python/aot/core:types",
         "//litert/python/aot/vendors",  # buildcleaner: keep
         "//litert/python/aot/vendors:import_vendor",
         "@tqdm",
@@ -78,10 +78,10 @@ pytype_strict_contrib_test(
     srcs = ["prepare_for_npu_test.py"],
     deps = [
         ":prepare_for_npu",
+        "//litert/python/aot/core:aot_types",
         "//litert/python/aot/core:common",
         "//litert/python/aot/core:components",
         "//litert/python/aot/core:test_common",
-        "//litert/python/aot/core:types",
         "//litert/python/aot/vendors/example:example_backend",
         "@absl_py//absl/testing:absltest",
     ],

litert/python/aot/ai_pack/BUILD

@@ -24,8 +24,8 @@ pytype_strict_library(
     srcs = ["export_lib.py"],
     visibility = ["//visibility:public"],
     deps = [
+        "//litert/python/aot/core:aot_types",
         "//litert/python/aot/core:common",
-        "//litert/python/aot/core:types",
         "//litert/python/aot/vendors:fallback_backend",
         "//litert/python/aot/vendors/google_tensor:target",
         "//litert/python/aot/vendors/mediatek:mediatek_backend",
@@ -40,7 +40,7 @@ pytype_strict_contrib_test(
     srcs = ["export_lib_test.py"],
     deps = [
         ":export_lib",
-        "//litert/python/aot/core:types",
+        "//litert/python/aot/core:aot_types",
         "//litert/python/aot/vendors:fallback_backend",
         "//litert/python/aot/vendors/google_tensor:target",
         "@absl_py//absl/testing:absltest",

litert/python/aot/ai_pack/export_lib.py

@@ -18,8 +18,8 @@
 import pathlib
 from typing import cast
 
+from litert.python.aot.core import aot_types
 from litert.python.aot.core import common
-from litert.python.aot.core import types
 from litert.python.aot.vendors import fallback_backend
 from litert.python.aot.vendors.google_tensor import target as google_tensor_target
 from litert.python.aot.vendors.mediatek import mediatek_backend
@@ -44,7 +44,7 @@
         </config:device-selector>"""
 
 
-def _is_mobile_device_backend(backend: types.Backend):
+def _is_mobile_device_backend(backend: aot_types.Backend):
   """Returns True if the backend is a mobile device backend."""
   target = backend.target
   if backend.id() == qualcomm_backend.QualcommBackend.id():
@@ -68,7 +68,7 @@ def _is_mobile_device_backend(backend: types.Backend):
 
 
 def _export_model_files_to_ai_pack(
-    compiled_models: types.CompilationResult,
+    compiled_models: aot_types.CompilationResult,
     ai_pack_dir: pathlib.Path,
     ai_pack_name: str,
     litert_model_name: str,
@@ -131,7 +131,7 @@ def _export_model_files_to_ai_pack(
 
 
 def _export_model_files_to_mtk_ai_pack(
-    compiled_models: types.CompilationResult,
+    compiled_models: aot_types.CompilationResult,
     ai_pack_dir: pathlib.Path,
     ai_pack_name: str,
     litert_model_name: str,
@@ -164,7 +164,7 @@ def _export_model_files_to_mtk_ai_pack(
     model.save(model_export_path, export_only=True)
 
 
-def _build_targeting_config(compiled_backends: list[types.Backend]) -> str:
+def _build_targeting_config(compiled_backends: list[aot_types.Backend]) -> str:
   """Builds device-targeting-config in device_targeting_configuration.xml."""
   device_groups = []
   for backend in compiled_backends:
@@ -178,7 +178,7 @@ def _build_targeting_config(compiled_backends: list[types.Backend]) -> str:
   return _DEVICE_TARGETING_CONFIGURATION.format(device_groups=device_groups)
 
 
-def _target_to_ai_pack_info(target: types.Target) -> str | None:
+def _target_to_ai_pack_info(target: aot_types.Target) -> str | None:
   """Builds the device group used in device_targeting_configuration.xml."""
   if isinstance(target, qnn_target.Target):
     group_name = str(target)
@@ -252,7 +252,7 @@ def _process_google_tensor_target(
 
 
 def _write_targeting_config(
-    compiled_models: types.CompilationResult, ai_pack_dir: pathlib.Path
+    compiled_models: aot_types.CompilationResult, ai_pack_dir: pathlib.Path
 ) -> None:
   """Writes device_targeting_configuration.xml for the given compiled models."""
   compiled_backends = [x for x, _ in compiled_models.models_with_backend]
@@ -265,7 +265,7 @@ def _write_targeting_config(
 
 
 def export(
-    compiled_models: types.CompilationResult,
+    compiled_models: aot_types.CompilationResult,
     ai_pack_dir: pathlib.Path | str,
     ai_pack_name: str,
     litert_model_name: str,

litert/python/aot/ai_pack/export_lib_test.py

@@ -17,28 +17,28 @@
 from unittest import mock
 from absl.testing import absltest as googletest
 from litert.python.aot.ai_pack import export_lib
-from litert.python.aot.core import types
+from litert.python.aot.core import aot_types
 from litert.python.aot.vendors import fallback_backend
 from litert.python.aot.vendors.google_tensor import target as google_tensor_target
 
 
 def _create_test_model():
-  test_model = types.Model(model_bytes=b"test_model")
+  test_model = aot_types.Model(model_bytes=b"test_model")
   return test_model
 
 
 def _create_mock_google_tensor_backend(
     soc_model: google_tensor_target.SocModel,
 ):
-  mock_backend = mock.Mock(spec=types.Backend)
+  mock_backend = mock.Mock(spec=aot_types.Backend)
   mock_backend.target = google_tensor_target.Target(soc_model=soc_model)
   mock_backend.target_id = str(mock_backend.target)
   mock_backend.id.return_value = google_tensor_target.Target.backend_id()
   return mock_backend
 
 
 def _create_mock_fallback_backend():
-  mock_backend = mock.Mock(spec=types.Backend)
+  mock_backend = mock.Mock(spec=aot_types.Backend)
   mock_backend.target = fallback_backend.FallbackTarget()
   mock_backend.target_id = fallback_backend.FallbackBackend.id()
   mock_backend.id.return_value = fallback_backend.FallbackBackend.id()
@@ -48,7 +48,7 @@ def _create_mock_fallback_backend():
 class ExportLibTest(googletest.TestCase):
 
   def test_export_with_fallback_model_succeeds(self):
-    mock_compilation_result = types.CompilationResult(
+    mock_compilation_result = aot_types.CompilationResult(
         models_with_backend=[
             (
                 _create_mock_fallback_backend(),
@@ -95,7 +95,7 @@ def test_export_with_fallback_model_succeeds(self):
       )
 
   def test_export_without_fallback_model_fails(self):
-    mock_compilation_result = types.CompilationResult()
+    mock_compilation_result = aot_types.CompilationResult()
     with self.assertRaises(AssertionError):
       export_lib.export(
           compiled_models=mock_compilation_result,
@@ -106,7 +106,7 @@ def test_export_without_fallback_model_fails(self):
 
   def test_export_with_google_tensor_target_succeeds(self):
     output_dir = self.create_tempdir().full_path
-    mock_compilation_result = types.CompilationResult(
+    mock_compilation_result = aot_types.CompilationResult(
         models_with_backend=[
             (
                 _create_mock_google_tensor_backend(

litert/python/aot/aot_compile.py

@@ -18,25 +18,25 @@
 import tempfile
 
 from litert.python.aot import prepare_for_npu as core
+from litert.python.aot.core import aot_types
 from litert.python.aot.core import apply_plugin
 from litert.python.aot.core import components
 from litert.python.aot.core import mlir_transforms
-from litert.python.aot.core import types
 from litert.python.aot.vendors import import_vendor
 
 
 def aot_compile(
-    input_model: types.Model | str,
+    input_model: aot_types.Model | str,
     output_dir: str | pathlib.Path | None = None,
-    target: types.Target | list[types.Target] | None = None,
+    target: aot_types.Target | list[aot_types.Target] | None = None,
     config: (
-        types.CompilationConfig | list[types.CompilationConfig] | None
+        aot_types.CompilationConfig | list[aot_types.CompilationConfig] | None
     ) = None,
     quantizer: components.AieQuantizerT | None = None,
     keep_going: bool = True,
     subgraphs_to_compile: list[int] | None = None,
     **kwargs,
-) -> types.CompilationResult:
+) -> aot_types.CompilationResult:
   """Prepares a TFLite model for NPU execution.
 
   High level command that erforms various backend specific pre-processing steps
@@ -66,16 +66,16 @@ def aot_compile(
   if config is None:
     if target is None:
       target = import_vendor.AllRegisteredTarget()
-    if isinstance(target, types.Target):
-      config = types.CompilationConfig(target=target)
+    if isinstance(target, aot_types.Target):
+      config = aot_types.CompilationConfig(target=target)
     elif isinstance(target, list):
-      config = [types.CompilationConfig(target=t) for t in target]
+      config = [aot_types.CompilationConfig(target=t) for t in target]
     else:
       raise ValueError("Unsupported target type.")
 
   if isinstance(input_model, str):
     input_path = pathlib.Path(input_model)
-    input_model = types.Model.create_from_path(input_path)
+    input_model = aot_types.Model.create_from_path(input_path)
 
   # Resolve output paths.
   temp_dir = None
@@ -94,10 +94,10 @@ def aot_compile(
   output_dir_path = pathlib.Path(output_dir)
   output_dir_path.mkdir(parents=True, exist_ok=True)
 
-  if isinstance(config, types.CompilationConfig) or not config:
+  if isinstance(config, aot_types.CompilationConfig) or not config:
     if config:
       # Make pytype happy.
-      assert isinstance(config, types.CompilationConfig)
+      assert isinstance(config, aot_types.CompilationConfig)
       kw_config = config.to_dict() | kwargs
     else:
       kw_config = kwargs

litert/python/aot/aot_compile_test.py

@@ -15,8 +15,8 @@
 
 from absl.testing import absltest as googletest
 from litert.python.aot import aot_compile as aot_compile_lib
+from litert.python.aot.core import aot_types
 from litert.python.aot.core import test_common
-from litert.python.aot.core import types
 from litert.python.aot.vendors.mediatek import target as mtk_target
 from litert.python.aot.vendors.qualcomm import target as qnn_target
 
@@ -31,7 +31,7 @@ def test_compile(self):
     for path in self.get_model_paths():
       results.append(
           aot_compile_lib.aot_compile(
-              types.Model.create_from_path(path),
+              aot_types.Model.create_from_path(path),
               output_dir=self.output_dir(),
               target=[sm8450_target, mt6989_target],
               keep_going=True,

litert/python/aot/core/BUILD

@@ -51,8 +51,8 @@ pytype_strict_library(
 )
 
 pytype_strict_library(
-    name = "types",
-    srcs = ["types.py"],
+    name = "aot_types",
+    srcs = ["aot_types.py"],
     visibility = ["//visibility:public"],
 )
 
@@ -66,9 +66,9 @@ pytype_strict_library(
         "//conditions:default": [],
     }),
     deps = [
+        ":aot_types",
         ":common",
         ":components",
-        ":types",
     ],
 )
 
@@ -81,9 +81,9 @@ pytype_strict_contrib_test(
         "nosan",
     ],
     deps = [
+        ":aot_types",
         ":apply_plugin",
         ":test_common",
-        ":types",
         "@absl_py//absl/testing:absltest",
     ],
 )
@@ -104,35 +104,35 @@ pytype_strict_library(
     name = "mlir_transforms",
     srcs = ["mlir_transforms.py"],
     deps = [
+        ":aot_types",
         ":components",
         ":tflxx_util",
-        ":types",
     ],
 )
 
 pytype_strict_library(
     name = "components",
     srcs = ["components.py"],
     visibility = ["//visibility:public"],
-    deps = [":types"],
+    deps = [":aot_types"],
 )
 
 pytype_strict_library(
     name = "tflxx_util",
     srcs = ["tflxx_util.py"],
     deps = [
-        ":types",
+        ":aot_types",
     ] + tflxx_deps_if_enabled(),
 )
 
 pytype_strict_contrib_test(
     name = "mlir_transforms_test",
     srcs = ["mlir_transforms_test.py"],
     deps = [
+        ":aot_types",
         ":common",
         ":mlir_transforms",
         ":test_common",
-        ":types",
         "@absl_py//absl/testing:absltest",
     ],
 )

litert/python/aot/core/types.py litert/python/aot/core/aot_types.pylitert/python/aot/core/types.py renamed to litert/python/aot/core/aot_types.py

@@ -21,9 +21,9 @@
 import subprocess
 import tempfile
 
+from litert.python.aot.core import aot_types
 from litert.python.aot.core import common
 from litert.python.aot.core import components
-from litert.python.aot.core import types
 
 _BINARY = pathlib.Path("tools/apply_plugin_main")
 
@@ -55,8 +55,8 @@ def component_name(self) -> str:
 
   def __call__(
       self,
-      input_model: types.Model,
-      output_model: types.Model,
+      input_model: aot_types.Model,
+      output_model: aot_types.Model,
       soc_manufacturer: str,
       soc_model: str,
       sdk_libs_path: str | None = None,
@@ -138,9 +138,9 @@ def __call__(
       raise ValueError(f"{output_model.path} is not a TFLite model.")
 
     partition_stats = _RE_PARTITION_STATS.findall(result.stdout)
-    output_model.partition_stats = types.PartitionStats(
+    output_model.partition_stats = aot_types.PartitionStats(
         subgraph_stats=[
-            types.SubgraphPartitionStats(
+            aot_types.SubgraphPartitionStats(
                 subgraph_index=int(s[0]),
                 num_ops_offloaded=int(s[1]),
                 num_total_ops=int(s[2]),