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

ai_edge_torch/lowertools/common_utils.py

@@ -0,0 +1,89 @@
+# 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.
+# ==============================================================================
+
+import logging
+
+from ai_edge_torch._convert import signature as signature_module
+import tensorflow as tf
+import torch
+
+
+def _torch_to_tf_tensor(torch_tensor: torch.Tensor):
+  if not torch_tensor.is_contiguous():
+    torch_tensor = torch_tensor.contiguous()
+
+  try:
+    dlpack_capsule = torch.utils.dlpack.to_dlpack(torch_tensor)
+    tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_capsule)
+  except Exception:
+    logging.info(
+        "Can not use dlpack to convert torch tensors. Falling back to numpy."
+    )
+    nparray = torch_tensor.cpu().detach().numpy()
+    tf_tensor = tf.convert_to_tensor(nparray)
+
+  return tf_tensor
+
+
+def _get_states(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature_module.Signature],
+):
+  for exported_program, signature in zip(exported_programs, signatures):
+    args, _ = exported_program.example_inputs
+    # Calling this to get **all** the state including model buffers.
+    _flat_input_args = exported_program._graph_module_flat_inputs(args, {})
+    for tensor, input_spec in zip(
+        _flat_input_args, exported_program.graph_signature.input_specs
+    ):
+      # Only interested in Tensors that are part of the state (and not user input).
+      if (
+          not isinstance(tensor, torch.Tensor)
+          or input_spec.kind
+          == torch.export.graph_signature.InputKind.USER_INPUT
+      ):
+        continue
+      yield signature, tensor, input_spec
+
+
+def _tensor_unique_id(tensor: torch.Tensor):
+  return (
+      str(tensor.device),
+      tensor.shape,
+      tensor.stride(),
+      tensor.untyped_storage().data_ptr(),
+  )
+
+
+def gather_state_dict(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature_module.Signature],
+):
+  deduped_tensor_map = {}
+
+  for _, tensor, _ in _get_states(exported_programs, signatures):
+    unique_id = _tensor_unique_id(tensor)
+    deduped_tensor_map[unique_id] = _torch_to_tf_tensor(tensor)
+
+  state_dict = {}
+  for signature, tensor, input_spec in _get_states(
+      exported_programs, signatures
+  ):
+    unique_id = _tensor_unique_id(tensor)
+    state_dict[signature.name + "_" + input_spec.target] = deduped_tensor_map[
+        unique_id
+    ]
+
+  return state_dict

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -0,0 +1,201 @@
+# 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.
+# ==============================================================================
+
+import tempfile
+from typing import Any, Dict, List, Optional, Tuple
+
+from ai_edge_torch import odml_torch
+from ai_edge_torch._convert import conversion_utils
+from ai_edge_torch._convert import signature as signature_module
+from ai_edge_torch.lowertools import common_utils
+from ai_edge_torch.odml_torch import export
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.quantize import quant_config as qcfg
+import tensorflow as tf
+import torch
+
+from tensorflow.compiler.tf2xla.python import xla as tfxla
+
+MlirBundle = odml_torch.export.MlirLowered
+MergedBundle = list[odml_torch.export.MlirLowered]
+
+
+def torch_dtype_to_tf(dtype):
+  return {
+      torch.double: tf.float64,
+      torch.float32: tf.float32,
+      torch.half: tf.float16,
+      torch.long: tf.int64,
+      torch.int32: tf.int32,
+      torch.int16: tf.int16,
+      torch.bool: tf.bool,
+  }.get(dtype)
+
+
+def _get_shape_with_dynamic(signature: export.VariableSignature):
+  return [
+      None if export_utils.is_torch_dynamic(s) else s for s in signature.shape
+  ]
+
+
+def _extract_call_args(
+    bundle: export.MlirLowered,
+    args: Tuple[Any],
+    tf_state_dict: Dict[str, tf.Variable],
+):
+  call_args = []
+  for sig in bundle.input_signature:
+    if sig.input_spec.is_user_input:
+      call_args.append(args[sig.input_spec.i])
+    elif sig.input_spec.is_parameter:
+      name = sig.input_spec.name
+      call_args.append(tf_state_dict[name])
+  return call_args
+
+
+def _wrap_as_tf_func(bundle, tf_state_dict):
+  def inner(*args):
+    t_outs = [torch_dtype_to_tf(sig.dtype) for sig in bundle.output_signature]
+    s_outs = [_get_shape_with_dynamic(sig) for sig in bundle.output_signature]
+    call_args = _extract_call_args(bundle, args, tf_state_dict)
+    return tfxla.call_module(
+        tuple(call_args),
+        version=5,
+        Tout=t_outs,  # dtype information
+        Sout=s_outs,  # Shape information
+        function_list=[],
+        module=bundle.module_bytecode,
+    )
+
+  return inner
+
+
+def _make_tf_signature(
+    input_signature: list[export.VariableSignature],
+    signature: signature_module.Signature,
+) -> List[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
+  user_input_signature = sorted(
+      [sig for sig in input_signature if sig.input_spec.is_user_input],
+      key=lambda sig: sig.input_spec.i,
+  )
+  tf_signature = []
+
+  for sig in user_input_signature:
+    shape = _get_shape_with_dynamic(sig)
+    tf_signature.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=torch_dtype_to_tf(sig.dtype),
+            name=input_names[sig.input_spec.i],
+        )
+    )
+  return tf_signature
+
+
+def merged_bundle_to_tfl_model(
+    merged_bundle: MergedBundle,
+    signatures: list[signature_module.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: dict = {},
+):
+  tf_state_dict = {
+      k: tf.Variable(v, trainable=False)
+      for k, v in merged_bundle[0].state_dict.items()
+  }
+
+  tf_signatures = [
+      _make_tf_signature(bundle.input_signature, sig)
+      for bundle, sig in zip(merged_bundle, signatures)
+  ]
+  tf_functions = [
+      _wrap_as_tf_func(bundle, tf_state_dict) for bundle in merged_bundle
+  ]
+
+  tf_module = tf.Module()
+  tf_module.f = []
+
+  for tf_sig, func in zip(tf_signatures, tf_functions):
+    tf_module.f.append(
+        tf.function(
+            func,
+            input_signature=tf_sig,
+        )
+    )
+
+  tf_module._variables = list(tf_state_dict.values())
+
+  tf_concrete_funcs = [
+      func.get_concrete_function(*tf_sig)
+      for func, tf_sig in zip(tf_module.f, tf_signatures)
+  ]
+
+  # We need to temporarily save since TFLite's from_concrete_functions does not
+  # allow providing names for each of the concrete functions.
+  with tempfile.TemporaryDirectory() as temp_dir_path:
+    tf.saved_model.save(
+        tf_module,
+        temp_dir_path,
+        signatures={
+            sig.name: tf_concrete_funcs[idx]
+            for idx, sig in enumerate(signatures)
+        },
+    )
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    converter._experimental_enable_composite_direct_lowering = True
+
+    conversion_utils.apply_tfl_converter_flags(converter, _tfl_converter_flags)
+
+    tflite_model = converter.convert()
+
+  return tflite_model
+
+
+def exported_program_to_mlir_text(
+    exported_program: torch.export.ExportedProgram,
+) -> str:
+  """Converts a ExportedProgram to a MLIR text."""
+  return odml_torch.export.exported_program_to_mlir(exported_program).get_text(
+      enable_debug_info=True
+  )
+
+
+def exported_program_to_mlir(
+    exported_program: torch.export.ExportedProgram,
+    sample_args: tuple[torch.Tensor],
+) -> export.MlirLowered:
+  """Converts a ExportedProgram to a MlirLowered."""
+  return odml_torch.export.exported_program_to_mlir(exported_program)
+
+
+def merge_mlir_bundles(
+    bundles: list[export.MlirLowered],
+    signatures: list[signature_module.Signature],
+    exported_programs: list[torch.export.ExportedProgram],
+) -> MergedBundle:
+  """Merges a list of MlirLowered into one."""
+  state_dict = common_utils.gather_state_dict(exported_programs, signatures)
+
+  merged_bundle = bundles.copy()
+  for bundle, signature in zip(merged_bundle, signatures):
+    bundle.state_dict = state_dict
+
+    for var_sig in bundle.input_signature:
+      if var_sig.input_spec.is_parameter:
+        var_sig.input_spec.name = signature.name + "_" + var_sig.input_spec.name
+
+  return merged_bundle

ai_edge_torch/{convert/conversion_utils.py → lowertools/torch_xla_utils.py}

@@ -13,19 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 
-import collections
 import copy
 from dataclasses import dataclass
 import gc
 import itertools
 import logging
 import tempfile
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
+from ai_edge_torch import model
+from ai_edge_torch._convert import conversion_utils
+from ai_edge_torch._convert import signature as signature_module
 from ai_edge_torch.generative.quantize.ai_edge_quantizer_glue import translate_recipe  # NOQA
+from ai_edge_torch.lowertools import common_utils
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
-import torch.utils._pytree as pytree
 from torch_xla import stablehlo
 
 try:
@@ -41,92 +43,11 @@ except ImportError:
   )
   raise
 
-DEFAULT_SIGNATURE_NAME = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
-
-
-@dataclass
-class Signature:
-  name: str
-  module: torch.nn.Module
-  sample_args: tuple[torch.Tensor]
-  sample_kwargs: dict[str, torch.Tensor]
-  dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None
-
-  @property
-  def _normalized_sample_args_kwargs(self):
-    args, kwargs = self.sample_args, self.sample_kwargs
-    if args is not None:
-      if not isinstance(args, tuple):
-        # TODO(b/352584188): Check value types
-        raise ValueError("sample_args must be a tuple of torch tensors.")
-    if kwargs is not None:
-      if not isinstance(kwargs, dict) or not all(
-          isinstance(key, str) for key in kwargs.keys()
-      ):
-        # TODO(b/352584188): Check value types
-        raise ValueError("sample_kwargs must be a dict of string to tensor.")
-
-    args = args if args is not None else tuple()
-    kwargs = kwargs if kwargs is not None else {}
-    return args, kwargs
-
-  @property
-  def flat_arg_names(self) -> list[str]:
-    spec = pytree.tree_flatten(self._normalized_sample_args_kwargs)[1]
-    args_spec, kwargs_spec = spec.children_specs
-
-    names = []
-    for i in range(args_spec.num_leaves):
-      names.append(f"args_{i}")
-
-    kwargs_names = self._flat_kwarg_names(
-        kwargs_spec.children_specs, kwargs_spec.context
-    )
-    names.extend(kwargs_names)
-    return names
-
-  def _flat_kwarg_names(self, specs, context) -> List[str]:
-    flat_names = []
-    if context is None:
-      for i, spec in enumerate(specs):
-        if spec.children_specs:
-          flat_names.extend([
-              f"{i}_{name}"
-              for name in self._flat_kwarg_names(
-                  spec.children_specs, spec.context
-              )
-          ])
-        else:
-          flat_names.append(f"{i}")
-    else:
-      flat_ctx = self._flatten_list(context)
-      for prefix, spec in zip(flat_ctx, specs):
-        leaf_flat_names = self._flat_kwarg_names(
-            spec.children_specs, spec.context
-        )
-        if leaf_flat_names:
-          flat_names.extend([f"{prefix}_{name}" for name in leaf_flat_names])
-        else:
-          flat_names.append(prefix)
-
-    return flat_names
-
-  def _flatten_list(self, l: List) -> List:
-    flattened = []
-    for item in l:
-      if isinstance(item, list):
-        flattened.extend(self._flatten_list(item))
-      else:
-        flattened.append(item)
-    return flattened
-
-  @property
-  def flat_args(self) -> tuple[Any]:
-    args, kwargs = self._normalized_sample_args_kwargs
-    return tuple([*args, *kwargs.values()])
+MlirBundle = stablehlo.StableHLOModelBundle
+MergedBundle = stablehlo.StableHLOModelBundle
 
 
-def exported_program_to_stablehlo_bundle(
+def exported_program_to_mlir(
     exported_program: torch.export.ExportedProgram,
     sample_args: tuple[torch.Tensor],
 ) -> stablehlo.StableHLOModelBundle:
@@ -141,81 +62,12 @@ def exported_program_to_stablehlo_bundle(
   )._bundle
 
 
-def _torch_to_tf_tensor(torch_tensor: torch.Tensor):
-  if not torch_tensor.is_contiguous():
-    torch_tensor = torch_tensor.contiguous()
-
-  try:
-    dlpack_capsule = torch.utils.dlpack.to_dlpack(torch_tensor)
-    tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_capsule)
-  except Exception:
-    logging.info(
-        "Can not use dlpack to convert torch tensors. Falling back to numpy."
-    )
-    nparray = torch_tensor.cpu().detach().numpy()
-    tf_tensor = tf.convert_to_tensor(nparray)
-
-  return tf_tensor
-
-
-def _get_states(
-    exported_programs: list[torch.export.ExportedProgram],
-    signatures: list[Signature],
-):
-  for exported_program, signature in zip(exported_programs, signatures):
-    args, _ = exported_program.example_inputs
-    # Calling this to get **all** the state including model buffers.
-    _flat_input_args = exported_program._graph_module_flat_inputs(args, {})
-    for tensor, input_spec in zip(
-        _flat_input_args, exported_program.graph_signature.input_specs
-    ):
-      # Only interested in Tensors that are part of the state (and not user input).
-      if (
-          not isinstance(tensor, torch.Tensor)
-          or input_spec.kind
-          == torch.export.graph_signature.InputKind.USER_INPUT
-      ):
-        continue
-      yield signature, tensor, input_spec
-
-
-def _tensor_unique_id(tensor: torch.Tensor):
-  return (
-      str(tensor.device),
-      tensor.shape,
-      tensor.stride(),
-      tensor.untyped_storage().data_ptr(),
-  )
-
-
-def _gather_state_dict(
-    exported_programs: list[torch.export.ExportedProgram],
-    signatures: list[Signature],
-):
-  deduped_tensor_map = {}
-
-  for _, tensor, _ in _get_states(exported_programs, signatures):
-    unique_id = _tensor_unique_id(tensor)
-    deduped_tensor_map[unique_id] = _torch_to_tf_tensor(tensor)
-
-  state_dict = {}
-  for signature, tensor, input_spec in _get_states(
-      exported_programs, signatures
-  ):
-    unique_id = _tensor_unique_id(tensor)
-    state_dict[signature.name + "_" + input_spec.target] = deduped_tensor_map[
-        unique_id
-    ]
-
-  return state_dict
-
-
-def merge_stablehlo_bundles(
+def merge_mlir_bundles(
     bundles: list[stablehlo.StableHLOModelBundle],
-    signatures: list[Signature],
+    signatures: list[signature_module.Signature],
     exported_programs: list[torch.export.ExportedProgram],
 ) -> stablehlo.StableHLOGraphModule:
-  state_dict = _gather_state_dict(exported_programs, signatures)
+  state_dict = common_utils.gather_state_dict(exported_programs, signatures)
 
   new_bundle = stablehlo.StableHLOModelBundle(
       state_dict=state_dict, additional_constants=[], stablehlo_funcs=[]
@@ -232,7 +84,7 @@ def merge_stablehlo_bundles(
           loc.name = signature.name + "_" + loc.name
       new_bundle.stablehlo_funcs.append(func)
     new_bundle.additional_constants.extend(bundle.additional_constants)
-  return stablehlo.StableHLOGraphModule(new_bundle)
+  return new_bundle
 
 
 def _get_shape_with_dynamic(signature: stablehlo.VariableSignature):
@@ -264,19 +116,15 @@ def _wrap_as_tf_func(
 
 
 def _make_tf_function(
-    shlo_graph_module: stablehlo.StableHLOGraphModule,
     bundle: stablehlo.StableHLOModelBundle = None,
 ):
-  bundle = shlo_graph_module._bundle if bundle is None else bundle
-  return [
-      _wrap_as_tf_func(func, bundle)
-      for func in shlo_graph_module._bundle.stablehlo_funcs
-  ]
+  bundle = bundle if bundle is None else bundle
+  return [_wrap_as_tf_func(func, bundle) for func in bundle.stablehlo_funcs]
 
 
 def _make_tf_signature(
     meta: stablehlo.StableHLOFunctionMeta,
-    signature: Signature,
+    signature: signature_module.Signature,
 ) -> list[tf.TensorSpec]:
   input_names = signature.flat_arg_names
   input_pos_to_spec = {
@@ -305,60 +153,33 @@ def _make_tf_signature(
   return ret
 
 
-def _apply_tfl_backdoor_flags(
-    converter: tf.lite.TFLiteConverter, tfl_converter_flags: dict
-):
-  def _set_converter_flag(path: list):
-    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, path: list):
-    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"
+def exported_program_to_mlir_text(
+    exported_program: torch.export.ExportedProgram,
+) -> str:
+  """Converts a ExportedProgram to a MLIR text."""
+  return stablehlo.exported_program_to_stablehlo(
+      exported_program
+  ).get_stablehlo_text()
 
 
-def convert_stablehlo_to_tflite(
-    shlo_graph_module: stablehlo.StableHLOGraphModule,
-    signatures: list[Signature],
+def merged_bundle_to_tfl_model(
+    bundle: stablehlo.StableHLOModelBundle,
+    signatures: list[signature_module.Signature],
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: dict = {},
 ) -> None:
   """Converts a StableHLOGraphModule to a tflite model.
-  Args:
-    shlo_graph_module - model to export and save
-    signatures: List of signatures from which names of the signatures is extracted.
+
+  Args: shlo_bundle - model to export and save
+
+    signatures: List of signatures from which names of the signatures is
+    extracted.
     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.
   """
 
-  bundle = shlo_graph_module._bundle
   tf_module = tf.Module()
   bundle.state_dict = {
       k: tf.Variable(v, trainable=False) for k, v in bundle.state_dict.items()
@@ -371,7 +192,7 @@ def convert_stablehlo_to_tflite(
       for func, sig in zip(bundle.stablehlo_funcs, signatures)
   )
 
-  tf_functions = _make_tf_function(shlo_graph_module, bundle)
+  tf_functions = _make_tf_function(bundle)
 
   tf_module.f = []
   for tf_sig, func in zip(tf_signatures, tf_functions):
@@ -413,7 +234,7 @@ def convert_stablehlo_to_tflite(
     converter._set_original_model_type(conversion_metadata_fb.ModelType.PYTORCH)
     converter._experimental_enable_composite_direct_lowering = True
 
-    _set_tfl_converter_quant_flags(converter, quant_config)
+    conversion_utils.set_tfl_converter_quant_flags(converter, quant_config)
     if (
         quant_config is not None
         and quant_config._quantizer_mode
@@ -423,7 +244,7 @@ def convert_stablehlo_to_tflite(
           quant_config.generative_recipe
       )
 
-    _apply_tfl_backdoor_flags(converter, _tfl_converter_flags)
+    conversion_utils.apply_tfl_converter_flags(converter, _tfl_converter_flags)
 
     tflite_model = converter.convert()
 

ai_edge_torch/model.py

@@ -15,17 +15,18 @@
 
 """Represents an ai_edge_torch model.
 
-PyTorch models can be converted to this representation through `ai_edge_torch.convert`.
+PyTorch models can be converted to this representation through
+`ai_edge_torch.convert`.
 """
 from __future__ import annotations
 
 import abc
 
-from ai_edge_torch.convert import conversion_utils as cutils
-import numpy as np
 import numpy.typing as npt
 import tensorflow as tf
 
+DEFAULT_SIGNATURE_NAME = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+
 
 class Model(abc.ABC):
   """Represents and edge model."""
@@ -34,7 +35,7 @@ class Model(abc.ABC):
   def __call__(
       self,
       *args: npt.ArrayLike,
-      signature_name: str = cutils.DEFAULT_SIGNATURE_NAME,
+      signature_name: str = DEFAULT_SIGNATURE_NAME,
       **kwargs,
   ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
     raise NotImplementedError()
@@ -66,18 +67,22 @@ class TfLiteModel(Model):
   def __call__(
       self,
       *args: npt.ArrayLike,
-      signature_name: str = cutils.DEFAULT_SIGNATURE_NAME,
+      signature_name: str = DEFAULT_SIGNATURE_NAME,
       **kwargs,
   ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
     """Runs inference on the edge model using the provided arguments.
 
     Args:
       *args: The arguments to be passed to the model for inference.
-      **kwargs: The arguments with specific names to be passed to the model for inference.
-      signature_name: The name of the signature to be used for inference.
-        The default signature is used if not provided.
+      **kwargs: The arguments with specific names to be passed to the model for
+        inference.
+      signature_name: The name of the signature to be used for inference. The
+        default signature is used if not provided.
     """
-    interpreter = tf.lite.Interpreter(model_content=self._tflite_model)
+    interpreter = tf.lite.Interpreter(
+        model_content=self._tflite_model,
+        experimental_default_delegate_latest_features=True,
+    )
     interpreter.allocate_tensors()
 
     signature_list = interpreter.get_signature_list()