Add a shim layer to allow usage of odml_torch (google-ai-edge#121)

Co-authored-by: Google AI Edge <[email protected]>

Author: majiddadashi

.github/workflows/formatting.yml

@@ -29,7 +29,7 @@ jobs:
           ref: ${{ inputs.trigger-sha }}
       - name: Install dependencies
         run: |
-          pip3 install pyink isort
+          pip3 install pyink
       - name: Check code style
         run: |
           ci/test_code_style.sh

WORKSPACE

@@ -46,14 +46,14 @@ http_archive(
     urls = [
         "https://github.com/google/sentencepiece/archive/refs/tags/v0.1.96.zip"
     ],
-    build_file = "@//third_party:sentencepiece.BUILD",
-    patches = ["@//third_party:com_google_sentencepiece.diff"],
+    build_file = "@//bazel:sentencepiece.BUILD",
+    patches = ["@//bazel:com_google_sentencepiece.diff"],
     patch_args = ["-p1"],
 )
 
 http_archive(
     name = "darts_clone",
-    build_file = "@//third_party:darts_clone.BUILD",
+    build_file = "@//bazel:darts_clone.BUILD",
     sha256 = "c97f55d05c98da6fcaf7f9ecc6a6dc6bc5b18b8564465f77abff8879d446491c",
     strip_prefix = "darts-clone-e40ce4627526985a7767444b6ed6893ab6ff8983",
     urls = [
@@ -92,7 +92,7 @@ http_archive(
       "https://github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
     ],
     patches = [
-        "@//third_party:org_tensorflow_system_python.diff",
+        "@//bazel:org_tensorflow_system_python.diff",
     ],
     patch_args = [
         "-p1",

ai_edge_torch/README.md

@@ -1,4 +1,3 @@
 
 * Documentation of the [PyTorch Converter](../docs/pytorch_converter/README.md)
 * Documentation of the [Generative API](generative/)
-

ai_edge_torch/__init__.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-from .convert.converter import convert
-from .convert.converter import signature
-from .convert.to_channel_last_io import to_channel_last_io
-from .model import Model
-from .version import __version__
+from ai_edge_torch._convert.converter import convert
+from ai_edge_torch._convert.converter import signature
+from ai_edge_torch._convert.to_channel_last_io import to_channel_last_io
+from ai_edge_torch.model import Model
+from ai_edge_torch.version import __version__
 
 
 def load(path: str) -> Model:

ai_edge_torch/convert/__init__.py renamed to ai_edge_torch/_convert/__init__.py

@@ -13,48 +13,44 @@
 # limitations under the License.
 # ==============================================================================
 
-import gc
 import logging
 import os
-from typing import Optional
+from typing import Any, Optional
 
+from ai_edge_torch import lowertools
 from ai_edge_torch import model
-from ai_edge_torch.convert import conversion_utils as cutils
-from ai_edge_torch.convert.fx_passes import BuildAtenCompositePass
-from ai_edge_torch.convert.fx_passes import BuildInterpolateCompositePass  # NOQA
-from ai_edge_torch.convert.fx_passes import CanonicalizePass
-from ai_edge_torch.convert.fx_passes import InjectMlirDebuginfoPass
-from ai_edge_torch.convert.fx_passes import OptimizeLayoutTransposesPass
-from ai_edge_torch.convert.fx_passes import run_passes
-from ai_edge_torch.generative.fx_passes import run_generative_passes
+from ai_edge_torch._convert import fx_passes
+from ai_edge_torch._convert import signature
+from ai_edge_torch.generative import fx_passes as generative_fx_passes
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
-from torch.export import ExportedProgram
-from torch_xla import stablehlo
 
 os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
 
 
 def _run_convert_passes(
-    exported_program: ExportedProgram,
-) -> ExportedProgram:
-  exported_program = run_generative_passes(exported_program)
-  return run_passes(
+    exported_program: torch.export.ExportedProgram,
+) -> torch.export.ExportedProgram:
+  exported_program = generative_fx_passes.run_generative_passes(
+      exported_program
+  )
+  return fx_passes.run_passes(
       exported_program,
       [
-          BuildInterpolateCompositePass(),
-          CanonicalizePass(),
-          OptimizeLayoutTransposesPass(),
-          CanonicalizePass(),
-          BuildAtenCompositePass(),
-          CanonicalizePass(),
-          InjectMlirDebuginfoPass(),
-          CanonicalizePass(),
+          fx_passes.BuildInterpolateCompositePass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.OptimizeLayoutTransposesPass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.BuildAtenCompositePass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.InjectMlirDebuginfoPass(),
+          fx_passes.CanonicalizePass(),
       ],
   )
 
 
-def _warn_training_modules(signatures: list[cutils.Signature]):
+def _warn_training_modules(signatures: list[signature.Signature]):
+  """Warns the user if the module is in training mode (.eval not called)."""
   for sig in signatures:
     if not sig.module.training:
       continue
@@ -64,30 +60,39 @@ def _warn_training_modules(signatures: list[cutils.Signature]):
         " module in evaluation mode with `module.eval()` for better on-device"
         " performance and compatibility."
     )
-    if len(signatures) == 1 and sig.name == cutils.DEFAULT_SIGNATURE_NAME:
+    if len(signatures) == 1 and sig.name == model.DEFAULT_SIGNATURE_NAME:
       # User does not specify any signature names explicitly.
       message = message.format(sig_name="")
     else:
       message = message.format(sig_name=f'"{sig.name}" ')
 
-    logging.warn(message)
+    logging.warning(message)
 
 
 def convert_signatures(
-    signatures: list[cutils.Signature],
+    signatures: list[signature.Signature],
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
-    _tfl_converter_flags: dict = {},
+    _tfl_converter_flags: Optional[dict[str, Any]],
 ) -> model.TfLiteModel:
-  """Converts a list of `Signature`s and embeds them into one `model.TfLiteModel`.
+  """Converts a list of `signature.Signature`s and embeds them into one `model.TfLiteModel`.
+
   Args:
-      signatures: The list of 'Signature' objects containing PyTorch modules to be converted.
+      signatures: The list of 'signature.Signature' objects containing PyTorch
+        modules to be converted.
       quant_config: User-defined quantization method and scheme of the model.
-      _tfl_converter_flags: A nested dictionary allowing setting flags for the underlying tflite converter.
+      _tfl_converter_flags: A nested dictionary allowing setting flags for the
+        underlying tflite converter.
+
+  Returns:
+    The converted `model.TfLiteModel` object.
   """
+  if _tfl_converter_flags is None:
+    _tfl_converter_flags = {}
+
   _warn_training_modules(signatures)
 
-  exported_programs: torch.export.ExportedProgram = [
+  exported_programs: torch.export.torch.export.ExportedProgram = [
       torch.export.export(
           sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes
       )
@@ -96,23 +101,8 @@ def convert_signatures(
 
   # Apply default fx passes
   exported_programs = list(map(_run_convert_passes, exported_programs))
-  shlo_bundles: list[stablehlo.StableHLOModelBundle] = [
-      cutils.exported_program_to_stablehlo_bundle(exported, sig.flat_args)
-      for exported, sig in zip(exported_programs, signatures)
-  ]
-
-  merged_shlo_graph_module: stablehlo.StableHLOGraphModule = (
-      cutils.merge_stablehlo_bundles(
-          shlo_bundles, signatures, exported_programs
-      )
-  )
-  del exported_programs
-  del shlo_bundles
-
-  gc.collect()
-
-  tflite_model = cutils.convert_stablehlo_to_tflite(
-      merged_shlo_graph_module,
+  tflite_model = lowertools.exported_programs_to_tflite(
+      exported_programs,
       signatures,
       quant_config=quant_config,
       _tfl_converter_flags=_tfl_converter_flags,

ai_edge_torch/convert/conversion.py renamed to ai_edge_torch/_convert/conversion.py

@@ -0,0 +1,64 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Any
+
+from ai_edge_torch.quantize import quant_config as qcfg
+import tensorflow as tf
+
+
+def apply_tfl_converter_flags(
+    converter: tf.lite.TFLiteConverter, tfl_converter_flags: dict[str, Any]
+):
+  """Applies TFLite converter flags to the converter.
+
+  Args:
+    converter: TFLite converter.
+    tfl_converter_flags: TFLite converter flags.
+  """
+
+  def _set_converter_flag(path: list[Any]):
+    if len(path) < 2:
+      raise ValueError("Expecting at least two values in the path.")
+
+    target_obj = converter
+    for idx in range(len(path) - 2):
+      target_obj = getattr(target_obj, path[idx])
+
+    setattr(target_obj, path[-2], path[-1])
+
+  def _iterate_dict_tree(flags_dict: dict[str, Any], path: list[Any]):
+    for key, value in flags_dict.items():
+      path.append(key)
+      if isinstance(value, dict):
+        _iterate_dict_tree(value, path)
+      else:
+        path.append(value)
+        _set_converter_flag(path)
+        path.pop()
+      path.pop()
+
+  _iterate_dict_tree(tfl_converter_flags, [])
+
+
+def set_tfl_converter_quant_flags(
+    converter: tf.lite.TFLiteConverter, quant_config: qcfg.QuantConfig
+):
+  if quant_config is not None:
+    quantizer_mode = quant_config._quantizer_mode
+    if quantizer_mode == qcfg.QuantConfig._QuantizerMode.PT2E_DYNAMIC:
+      converter._experimental_qdq_conversion_mode = "DYNAMIC"
+    elif quantizer_mode == qcfg.QuantConfig._QuantizerMode.PT2E_STATIC:
+      converter._experimental_qdq_conversion_mode = "STATIC"