ai-edge-torch-nightly 0.2.0.dev20240805__py3-none-any.whl → 0.2.0.dev20240808__py3-none-any.whl

ai_edge_torch/generative/utilities/stable_diffusion_loader.py

@@ -594,11 +594,13 @@ class AutoEncoderModelLoader(BaseLoader):
     up_decoder_blocks_tensor_names: List[UpDecoderBlockTensorNames] = None
 
   def __init__(self, file_name: str, names: TensorNames):
-    """AutoEncoderModelLoader constructor. Can be used to load encoder and decoder models.
+    """AutoEncoderModelLoader constructor.
+
+    Can be used to load encoder and decoder models.
 
     Args:
-        file_name (str): Path to the checkpoint. Can be a directory or an
-          exact file.
+        file_name (str): Path to the checkpoint. Can be a directory or an exact
+          file.
         names (TensorNames): An instance of `TensorNames` to determine mappings.
     """
     self._file_name = file_name
@@ -617,7 +619,8 @@ class AutoEncoderModelLoader(BaseLoader):
 
     Returns:
         missing_keys (List[str]): a list of str containing the missing keys.
-        unexpected_keys (List[str]): a list of str containing the unexpected keys.
+        unexpected_keys (List[str]): a list of str containing the unexpected
+        keys.
 
     Raises:
         ValueError: If conversion results in unmapped tensors and strict mode is
@@ -683,6 +686,31 @@ class AutoEncoderModelLoader(BaseLoader):
     return model.load_state_dict(converted_state, strict=strict)
 
 
+def build_attention_config(
+    num_heads,
+    dim,
+    num_query_groups,
+    rotary_percentage=0.0,
+    qkv_transpose_before_split=True,
+    qkv_use_bias=False,
+    output_proj_use_bias=True,
+    enable_kv_cache=False,
+    qkv_fused_interleaved=False,
+):
+
+  return layers_config.AttentionConfig(
+      num_heads=num_heads,
+      head_dim=dim // num_heads,
+      num_query_groups=num_query_groups,
+      rotary_percentage=rotary_percentage,
+      qkv_transpose_before_split=qkv_transpose_before_split,
+      qkv_use_bias=qkv_use_bias,
+      output_proj_use_bias=output_proj_use_bias,
+      enable_kv_cache=enable_kv_cache,
+      qkv_fused_interleaved=qkv_fused_interleaved,
+  )
+
+
 class DiffusionModelLoader(BaseLoader):
 
   @dataclass
@@ -696,11 +724,13 @@ class DiffusionModelLoader(BaseLoader):
     up_decoder_blocks_tensor_names: List[UpDecoderBlockTensorNames] = None
 
   def __init__(self, file_name: str, names: TensorNames):
-    """DiffusionModelLoader constructor. Can be used to load diffusion models of Stable Diffusion.
+    """DiffusionModelLoader constructor.
+
+    Can be used to load diffusion models of Stable Diffusion.
 
     Args:
-        file_name (str): Path to the checkpoint. Can be a directory or an
-          exact file.
+        file_name (str): Path to the checkpoint. Can be a directory or an exact
+          file.
         names (TensorNames): An instance of `TensorNames` to determine mappings.
     """
     self._file_name = file_name
@@ -719,7 +749,8 @@ class DiffusionModelLoader(BaseLoader):
 
     Returns:
         missing_keys (List[str]): a list of str containing the missing keys.
-        unexpected_keys (List[str]): a list of str containing the unexpected keys.
+        unexpected_keys (List[str]): a list of str containing the unexpected
+        keys.
 
     Raises:
         ValueError: If conversion results in unmapped tensors and strict mode is
@@ -741,16 +772,6 @@ class DiffusionModelLoader(BaseLoader):
         state, self._names.final_norm, converted_state, "final_norm"
     )
 
-    attention_config = layers_config.AttentionConfig(
-        num_heads=config.transformer_num_attention_heads,
-        num_query_groups=config.transformer_num_attention_heads,
-        rotary_percentage=0.0,
-        qkv_transpose_before_split=True,
-        qkv_use_bias=False,
-        output_proj_use_bias=True,
-        enable_kv_cache=False,
-    )
-
     # Map down_encoders.
     output_channel = config.block_out_channels[0]
     for i, block_out_channel in enumerate(config.block_out_channels):
@@ -781,13 +802,21 @@ class DiffusionModelLoader(BaseLoader):
                 attention_block_config=unet_config.AttentionBlock2DConfig(
                     dim=output_channel,
                     normalization_config=config.transformer_norm_config,
-                    attention_config=attention_config,
+                    attention_config=build_attention_config(
+                        num_heads=config.transformer_num_attention_heads,
+                        dim=output_channel,
+                        num_query_groups=config.transformer_num_attention_heads,
+                    ),
                 ),
                 cross_attention_block_config=unet_config.CrossAttentionBlock2DConfig(
                     query_dim=output_channel,
                     cross_dim=config.transformer_cross_attention_dim,
                     normalization_config=config.transformer_norm_config,
-                    attention_config=attention_config,
+                    attention_config=build_attention_config(
+                        num_heads=config.transformer_num_attention_heads,
+                        dim=output_channel,
+                        num_query_groups=config.transformer_num_attention_heads,
+                    ),
                 ),
                 pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                 feed_forward_block_config=unet_config.FeedForwardBlock2DConfig(
@@ -839,13 +868,21 @@ class DiffusionModelLoader(BaseLoader):
             attention_block_config=unet_config.AttentionBlock2DConfig(
                 dim=mid_block_channels,
                 normalization_config=config.transformer_norm_config,
-                attention_config=attention_config,
+                attention_config=build_attention_config(
+                    num_heads=config.transformer_num_attention_heads,
+                    dim=mid_block_channels,
+                    num_query_groups=config.transformer_num_attention_heads,
+                ),
             ),
             cross_attention_block_config=unet_config.CrossAttentionBlock2DConfig(
                 query_dim=mid_block_channels,
                 cross_dim=config.transformer_cross_attention_dim,
                 normalization_config=config.transformer_norm_config,
-                attention_config=attention_config,
+                attention_config=build_attention_config(
+                    num_heads=config.transformer_num_attention_heads,
+                    dim=mid_block_channels,
+                    num_query_groups=config.transformer_num_attention_heads,
+                ),
             ),
             pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
             feed_forward_block_config=unet_config.FeedForwardBlock2DConfig(
@@ -904,13 +941,21 @@ class DiffusionModelLoader(BaseLoader):
                 attention_block_config=unet_config.AttentionBlock2DConfig(
                     dim=output_channel,
                     normalization_config=config.transformer_norm_config,
-                    attention_config=attention_config,
+                    attention_config=build_attention_config(
+                        num_heads=config.transformer_num_attention_heads,
+                        dim=output_channel,
+                        num_query_groups=config.transformer_num_attention_heads,
+                    ),
                 ),
                 cross_attention_block_config=unet_config.CrossAttentionBlock2DConfig(
                     query_dim=output_channel,
                     cross_dim=config.transformer_cross_attention_dim,
                     normalization_config=config.transformer_norm_config,
-                    attention_config=attention_config,
+                    attention_config=build_attention_config(
+                        num_heads=config.transformer_num_attention_heads,
+                        dim=output_channel,
+                        num_query_groups=config.transformer_num_attention_heads,
+                    ),
                 ),
                 pre_conv_normalization_config=config.transformer_pre_conv_norm_config,
                 feed_forward_block_config=unet_config.FeedForwardBlock2DConfig(

ai_edge_torch/generative/utilities/t5_loader.py

@@ -92,9 +92,7 @@ def load_pytorch_statedict(full_path: str):
 
 
 class ModelLoader:
-  """A utility class for loading and converting model checkpoints to ODML
-  transformer layer format.
-  """
+  """Utility class for loading and converting checkpoints to ODML transformer layer format."""
 
   @dataclass
   class TensorNames:
@@ -121,12 +119,13 @@ class ModelLoader:
     lm_head: str = None
 
   def __init__(self, file_name: str, names: TensorNames) -> None:
-    """ModelLoader constructor. Can be used to load multiple models of the same
-    type.
+    """ModelLoader constructor.
+
+    Can be used to load multiple models of the same type.
 
     Args:
-        file_name (str): Path to the checkpoint. Can be a directory or an
-          exact file.
+        file_name (str): Path to the checkpoint. Can be a directory or an exact
+          file.
         names (TensorNames): An instance of `TensorNames` to determine mappings.
     """
     self._file_name = file_name
@@ -158,7 +157,7 @@ class ModelLoader:
       )
     elif isinstance(self._names, dict):
       converted_state = {}
-      for additional_prefix, names in self._names.items():
+      for additional_prefix, _ in self._names.items():
         local_converted_state = self._do_load(
             model,
             state,
@@ -212,7 +211,7 @@ class ModelLoader:
 
     if names.relative_attn_bias:
       rel_attn_name = names.relative_attn_bias
-      prefix = additional_prefix + f"transformer_blocks.0"
+      prefix = additional_prefix + "transformer_blocks.0"
       converted_state[f"{prefix}.atten_func.relative_attention_bias.weight"] = (
           state.pop(f"{rel_attn_name}.weight")
       )
@@ -266,7 +265,7 @@ class ModelLoader:
     if self._file_name.endswith(".bin"):
       return load_pytorch_statedict
 
-    raise ValueError(f"File format not supported.")
+    raise ValueError("File format not supported.")
 
   def _map_feedforward(
       self,
@@ -505,8 +504,8 @@ class ModelLoader:
     q_per_kv = (
         config.attn_config.num_heads // config.attn_config.num_query_groups
     )
-    qs = torch.split(q, config.head_dim * q_per_kv)
-    ks = torch.split(k, config.head_dim)
-    vs = torch.split(v, config.head_dim)
+    qs = torch.split(q, config.attn_config.head_dim * q_per_kv)
+    ks = torch.split(k, config.attn_config.head_dim)
+    vs = torch.split(v, config.attn_config.head_dim)
     cycled = [t for group in zip(qs, ks, vs) for t in group]
     return torch.cat(cycled)

ai_edge_torch/hlfb/__init__.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-from torch_xla.experimental.mark_pattern_utils import StableHLOCompositeBuilder
+from ai_edge_torch.lowertools import StableHLOCompositeBuilder

ai_edge_torch/hlfb/mark_pattern/__init__.py

@@ -16,10 +16,10 @@ import copy
 from typing import Any
 import uuid
 
-from ai_edge_torch.hlfb.mark_pattern.pattern import Pattern
-from ai_edge_torch.hlfb.mark_pattern.pattern import ScalarAttrTracker  # NOQA
+from ai_edge_torch import lowertools
+from ai_edge_torch.hlfb.mark_pattern import passes
+from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
-from torch_xla.experimental import xla_marker
 
 
 @torch._dynamo.assume_constant_result
@@ -48,10 +48,10 @@ def _insert_marker(
     is_input: bool,
     attr: dict[str, Any] = None,
 ):
-  attr = xla_marker.serialize_composite_attr(attr) if attr else None
+  attr = lowertools.serialize_composite_attr(attr) if attr else None
   with graph_module.graph.inserting_after(node):
     new_node = graph_module.graph.call_function(
-        torch.ops.xla.mark_tensor,
+        lowertools.mark_tensor_op,
         args=(node,),
         kwargs={
             "name": name,
@@ -68,13 +68,16 @@ def _insert_marker(
 
 def mark_pattern(
     graph_module: torch.fx.GraphModule,
-    pattern: Pattern,
+    pattern: pattern_module.Pattern,
 ) -> torch.fx.GraphModule:
   """Mark all existences of pattern graph in the GraphModule with fx pattern matching.
+
   The marked subgraphs will be lowered in StableHLO composite ops.
+
   Args:
     graph_module (torch.fx.GraphModule): GraphModule to be matched and marked.
     pattern (ai_edge_torch.hlfb.mark_pattern.Pattern): Pattern to match.
+
   Returns:
     The modified graph_module with additional marker ops in graph.
   """

ai_edge_torch/hlfb/mark_pattern/passes.py

@@ -12,13 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Passes to clean up the model graph for pattern matching."""
+
 import torch
 
 
 def remove_clone_ops(gm: torch.fx.GraphModule):
-  # torch export adds additional aten.clone nodes to produce contiguous in memory tensors
-  # depending on tensor sizes for runtime efficiency. However, these unpredictable clone
-  # nodes can break the pattern matching. Thus remove all clones in model and pattern graphs.
+  """Removes clone ops from the graph.
+
+  torch export adds additional aten.clone nodes to produce contiguous in memory
+  tensors depending on tensor sizes for runtime efficiency. However, these
+  unpredictable clone nodes can break the pattern matching. Thus remove all
+  clones in model and pattern graphs.
+
+  Args:
+    gm: The graph module to remove clone ops from.
+
+  Returns:
+    The graph module with clone ops removed.
+  """
   for node in gm.graph.nodes:
     if node.op == "call_function" and node.name.startswith("clone"):
       node.replace_all_uses_with(node.args[0])
@@ -30,6 +42,14 @@ def remove_clone_ops(gm: torch.fx.GraphModule):
 
 
 def remove_dangling_args(gm: torch.fx.GraphModule):
+  """Removes dangling args from the graph.
+
+  Args:
+    gm: The graph module to remove dangling args from.
+
+  Returns:
+    The graph module with dangling args removed.
+  """
   nodes_to_erase = []
   for node in gm.graph.nodes:
     if node.op == "placeholder" and len(node.users) == 0:

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import copy
+"""Mark pattern."""
+
 import dataclasses
 from typing import Any, Callable, Optional, Union
 
@@ -45,6 +46,7 @@ def _are_equal(x: Any, y: Any) -> bool:
 @dataclasses.dataclass
 class ScalarAttrTracker:
   """ScalarAttrTracker is used to track the occurrence of a pattern's
+
   scalar arg/attr in the pattern decomposed graph. Since a scalar attr
   to the pattern can be transformed and turned into a/some ops' scalar
   arg in the decomposed graph, it would be hard to programmatically get
@@ -57,11 +59,10 @@ class ScalarAttrTracker:
     pattern_arg_pos (int): the index of the attr to track in the pattern's
       export_args.
     transform (Callable): the transform function used when targeting the
-      occurrence of the attr value in the decomposed graph. An attr value
-      may be transformed during the decomposition and appear as a derived
-      value.
-    inverse_transform (Callable): the inverse transform function that maps
-      the transformed value back to the original attr value.
+      occurrence of the attr value in the decomposed graph. An attr value may be
+      transformed during the decomposition and appear as a derived value.
+    inverse_transform (Callable): the inverse transform function that maps the
+      transformed value back to the original attr value.
   """
 
   attr_name: str
@@ -74,6 +75,7 @@ class ScalarAttrTracker:
 
   def track(self, *sources):
     """Register magic values to track the (transformed) attr values in
+
     the pattern decomposed graph.
     """
     for source in sources:
@@ -158,24 +160,22 @@ class Pattern:
     """The PyTorch computation pattern to match against a model.
 
     Args:
-      name (str): the name of the pattern. It would be propagated to
-        the `name` attr in StableHLO composite ops for the matched
-        model subgraphs in the lowering.
+      name (str): the name of the pattern. It would be propagated to the `name`
+        attr in StableHLO composite ops for the matched model subgraphs in the
+        lowering.
       module (torch.nn.Module or Callable): the PyTorch computation.
-      export_args (tuple[Any]): the args used to export the pattern module
-        with torch.export.export. If export_args contains non-tensor
-        Python scalars, there must be a corresponding attr tracker
-        in `scalar_attr_trackers` for each scalar arg.
-      attr_builder (Callable[[Pattern, GraphModule, InternalMatch], Optional[dict[str, Any]]]):
-        the callable that produces the a scalar attrs dict, which would be
-        propagated to `attr` in StableHLO composite ops for the matched
-        model subgraphs in the lowering.
-      scalar_attr_trackers (list[ScalarAttrTracker]): the trackers
-        for scalar args in `export_args`, which are used to track
-        the attr occurrence(s) and retrieve their values from the
-        matched subgraph.
-      decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]]):
-        The decomposition table to be run on the pattern's exported program.
+      export_args (tuple[Any]): the args used to export the pattern module with
+        torch.export.export. If export_args contains non-tensor Python scalars,
+        there must be a corresponding attr tracker in `scalar_attr_trackers` for
+        each scalar arg. attr_builder (Callable[[Pattern, GraphModule,
+        InternalMatch], Optional[dict[str, Any]]]): the callable that produces
+        the a scalar attrs dict, which would be propagated to `attr` in
+        StableHLO composite ops for the matched model subgraphs in the lowering.
+      scalar_attr_trackers (list[ScalarAttrTracker]): the trackers for scalar
+        args in `export_args`, which are used to track the attr occurrence(s)
+        and retrieve their values from the matched subgraph.
+      decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]]): The
+        decomposition table to be run on the pattern's exported program.
     """
     if not isinstance(module, torch.nn.Module):
 

ai_edge_torch/hlfb/test/test_mark_pattern.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Tests for mark_pattern."""
 
-import unittest
-
+from ai_edge_torch import lowertools
 from ai_edge_torch.hlfb import mark_pattern
+from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
-import torch_xla
+
+from tensorflow.python.platform import googletest
 
 
 def _export_stablehlo_mlir(model, args=None):
@@ -25,11 +27,10 @@ def _export_stablehlo_mlir(model, args=None):
     ep = torch.export.export(model, args)
   else:
     ep = model
-  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
-  return stablehlo_gm.get_stablehlo_text()
+  return lowertools.exported_program_to_mlir_text(ep)
 
 
-class TestMarkPattern(unittest.TestCase):
+class TestMarkPattern(googletest.TestCase):
 
   def test_mark_pattern(self):
 
@@ -38,7 +39,7 @@ class TestMarkPattern(unittest.TestCase):
       def forward(self, x):
         return x * x + x + x
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.add",
         lambda a, b: a + b,
         export_args=(torch.rand(2, 2), torch.rand(2, 2)),
@@ -58,7 +59,7 @@ class TestMarkPattern(unittest.TestCase):
       def forward(self, x):
         return x * x * x + x - x * x + x
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.add",
         lambda a, b: a + b,
         export_args=(torch.rand(2, 2), torch.rand(2, 2)),
@@ -85,12 +86,12 @@ class TestMarkPattern(unittest.TestCase):
           r = torch.nn.LogSoftmax(dim=idx % 2)(r) * x
         return r
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.log_softmax",
         lambda x, dim: torch.nn.functional.log_softmax(x, dim=dim),
         export_args=(torch.rand(10, 10, 10), 1),
         scalar_attr_trackers=[
-            mark_pattern.ScalarAttrTracker("dim", pattern_arg_pos=1)
+            pattern_module.ScalarAttrTracker("dim", pattern_arg_pos=1)
             .track(0)
             .track(1)
             .track(2),
@@ -115,7 +116,7 @@ class TestMarkPattern(unittest.TestCase):
         z = z + y
         return z
 
-    pattern = mark_pattern.Pattern(
+    pattern = pattern_module.Pattern(
         "test.relu",
         lambda x: torch.ops.aten.relu(x),
         export_args=(torch.rand(2, 2),),
@@ -131,4 +132,4 @@ class TestMarkPattern(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py

@@ -12,22 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Tests for StableHLOCompositeBuilder."""
+
 import math
-import unittest
 
+from ai_edge_torch import lowertools
 from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 import torch.nn.functional as F
-import torch_xla
+
+from tensorflow.python.platform import googletest
 
 
 def _export_stablehlo_mlir(model, args):
   ep = torch.export.export(model, args)
-  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
-  return stablehlo_gm.get_stablehlo_text()
+  return lowertools.exported_program_to_mlir_text(ep)
 
 
-class TestStableHLOCompositeBuilder(unittest.TestCase):
+class TestStableHLOCompositeBuilder(googletest.TestCase):
 
   def test_build_composite(self):
     class SampleModel(torch.nn.Module):
@@ -273,4 +275,4 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
 
 
 if __name__ == "__main__":
-  unittest.main()
+  googletest.main()

ai_edge_torch/{convert/fx_passes/optimize_layout_transposes_pass → lowertools}/__init__.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass.pass_body import OptimizeLayoutTransposesPass  # NOQA
+from ._shim import *

ai_edge_torch/lowertools/_shim.py

@@ -0,0 +1,80 @@
+# 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, Optional
+
+from ai_edge_torch import config
+from ai_edge_torch._convert import signature
+from ai_edge_torch.quantize import quant_config as qcfg
+import torch
+
+# isort: off
+if config.Config.use_torch_xla:
+  from ai_edge_torch.lowertools import torch_xla_utils as utils
+  from ai_edge_torch.lowertools.torch_xla_utils import exported_program_to_mlir_text
+  from torch_xla.experimental.mark_pattern_utils import StableHLOCompositeBuilder
+  from torch_xla.experimental.xla_marker import serialize_composite_attr
+  # The following imports are needed to register the needed torch_xla ops.
+  import torch_xla.experimental.xla_marker
+  import torch_xla.experimental.xla_mlir_debuginfo
+
+  mark_tensor_op = torch.ops.xla.mark_tensor.default
+  write_mlir_debuginfo_op = torch.ops.xla.write_mlir_debuginfo.default
+else:
+  from ai_edge_torch.lowertools import odml_torch_utils as utils
+  from ai_edge_torch.lowertools.odml_torch_utils import exported_program_to_mlir_text
+  from ai_edge_torch.odml_torch.composite import StableHLOCompositeBuilder
+  from ai_edge_torch.odml_torch.composite.mark_tensor import serialize_composite_attr
+  from ai_edge_torch.odml_torch.composite.mark_tensor import mark_tensor_op
+  from ai_edge_torch.odml_torch.debuginfo import write_mlir_debuginfo_op
+# isort: on
+
+
+def exported_programs_to_tflite(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: Optional[dict[str, Any]] = None,
+):
+  """Converts a list of ExportedProgram to a TFLite model.
+
+  Args:
+    exported_programs: A list of ExportedProgram.
+    signatures: A list of Signature.
+    quant_config: A QuantConfig.
+    _tfl_converter_flags: A dict of flags for TFLiteConverter.
+
+  Returns:
+    A TFLite model.
+  """
+  if _tfl_converter_flags is None:
+    _tfl_converter_flags = {}
+
+  bundles: list[utils.MlirBundle] = [
+      utils.exported_program_to_mlir(exported, sig.flat_args)
+      for exported, sig in zip(exported_programs, signatures)
+  ]
+
+  merged_bundle: utils.MergedBundle = utils.merge_mlir_bundles(
+      bundles, signatures, exported_programs
+  )
+
+  return utils.merged_bundle_to_tfl_model(
+      merged_bundle,
+      signatures,
+      quant_config=quant_config,
+      _tfl_converter_flags=_tfl_converter_flags,
+  )