update to tf.lite to use ai-edge-litert

python/tflite_micro/BUILD

@@ -94,6 +94,7 @@ py_test(
     ],
     deps = [
         ":runtime",
+        requirement("ai-edge-litert"),
         requirement("numpy"),
         requirement("tensorflow"),
         "//tensorflow/lite/micro/examples/recipes:add_four_numbers",

python/tflite_micro/runtime_test.py

@@ -25,6 +25,7 @@
 import numpy as np
 import tensorflow as tf
 
+from ai_edge_litert import interpreter as litert_interpreter
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import test
 from tflite_micro.python.tflite_micro import runtime
@@ -199,10 +200,10 @@ def testCompareWithTFLite(self):
     tflm_interpreter = runtime.Interpreter.from_bytes(model_data)
 
     # TFLite interpreter
-    tflite_interpreter = tf.lite.Interpreter(
+    tflite_interpreter = litert_interpreter.Interpreter(
         model_content=model_data,
         experimental_op_resolver_type=\
-        tf.lite.experimental.OpResolverType.BUILTIN_REF)
+        litert_interpreter.OpResolverType.BUILTIN_REF)
     tflite_interpreter.allocate_tensors()
     tflite_output_details = tflite_interpreter.get_output_details()[0]
     tflite_input_details = tflite_interpreter.get_input_details()[0]

tensorflow/lite/micro/examples/hello_world/BUILD

@@ -53,6 +53,7 @@ py_binary(
         "@absl_py//absl:app",
         "@absl_py//absl/flags",
         "@absl_py//absl/logging",
+        requirement("ai-edge-litert"),
         requirement("numpy"),
         requirement("tensorflow"),
         "//python/tflite_micro:runtime",

tensorflow/lite/micro/examples/hello_world/evaluate.py

@@ -16,6 +16,7 @@
 import tensorflow as tf
 from absl import app
 from absl import flags
+from ai_edge_litert import interpreter as litert_interpreter
 import numpy as np
 import matplotlib.pyplot as plt
 from tensorflow.python.platform import resource_loader
@@ -92,9 +93,9 @@ def get_tflm_prediction(model_path, x_values):
 # returns the prediction of the interpreter.
 def get_tflite_prediction(model_path, x_values):
   # TFLite interpreter
-  tflite_interpreter = tf.lite.Interpreter(
+  tflite_interpreter = litert_interpreter.Interpreter(
       model_path=model_path,
-      experimental_op_resolver_type=tf.lite.experimental.OpResolverType.
+      experimental_op_resolver_type=litert_interpreter.OpResolverType.
       BUILTIN_REF,
   )
   tflite_interpreter.allocate_tensors()

tensorflow/lite/micro/examples/mnist_lstm/BUILD

@@ -6,6 +6,7 @@ py_binary(
     srcs = ["train.py"],
     srcs_version = "PY3",
     deps = [
+        requirement("ai-edge-litert"),
         requirement("numpy"),
         requirement("tensorflow"),
     ],

tensorflow/lite/micro/examples/mnist_lstm/evaluate_test.py

@@ -17,6 +17,7 @@
 import numpy as np
 import tensorflow as tf
 
+from ai_edge_litert import interpreter as litert_interpreter
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import resource_loader
 from tensorflow.python.platform import test
@@ -43,10 +44,10 @@ def testInputErrHandling(self):
       evaluate.predict_image(self.tflm_interpreter, wrong_size_image_path)
 
   def testCompareWithTFLite(self):
-    tflite_interpreter = tf.lite.Interpreter(
+    tflite_interpreter = litert_interpreter.Interpreter(
         model_path=self.model_path,
         experimental_op_resolver_type=\
-        tf.lite.experimental.OpResolverType.BUILTIN_REF)
+        litert_interpreter.OpResolverType.BUILTIN_REF)
     tflite_interpreter.allocate_tensors()
     tflite_output_details = tflite_interpreter.get_output_details()[0]
     tflite_input_details = tflite_interpreter.get_input_details()[0]

tensorflow/lite/micro/tools/BUILD

@@ -34,6 +34,7 @@ py_library(
     srcs_version = "PY3",
     visibility = ["//:__subpackages__"],
     deps = [
+        requirement("ai-edge-litert"),
         "//tensorflow/lite/python:schema_py",
     ],
 )
@@ -208,6 +209,7 @@ py_binary(
         ":model_transforms_utils",
         "@absl_py//absl:app",
         "@absl_py//absl/flags",
+        requirement("ai-edge-litert"),
         requirement("tensorflow"),
         "//python/tflite_micro:runtime",
         "//tensorflow/lite/tools:flatbuffer_utils",

tensorflow/lite/micro/tools/generate_test_for_model.py

@@ -18,6 +18,7 @@
 import numpy as np
 import tensorflow as tf
 
+from ai_edge_litert import interpreter as litert_interpreter
 from tflite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
 
 
@@ -103,9 +104,9 @@ def generate_golden_single_in_single_out(self):
     if (len(self.model_paths) != 1):
       raise RuntimeError(f'Single model expected')
     model_path = self.model_paths[0]
-    interpreter = tf.lite.Interpreter(model_path=model_path,
+    interpreter = litert_interpreter.Interpreter(model_path=model_path,
                                       experimental_op_resolver_type=\
-                                      tf.lite.experimental.OpResolverType.BUILTIN_REF)
+                                      litert_interpreter.OpResolverType.BUILTIN_REF)
 
     interpreter.allocate_tensors()
 
@@ -140,10 +141,10 @@ def generate_goldens(self, builtin_operator):
 
     for model_path in self.model_paths:
       # Load model and run a single inference with random inputs.
-      interpreter = tf.lite.Interpreter(
+      interpreter = litert_interpreter.Interpreter(
           model_path=model_path,
           experimental_op_resolver_type=\
-          tf.lite.experimental.OpResolverType.BUILTIN_REF)
+          litert_interpreter.OpResolverType.BUILTIN_REF)
       interpreter.allocate_tensors()
       input_tensor = interpreter.tensor(
           interpreter.get_input_details()[0]['index'])

tensorflow/lite/micro/tools/layer_by_layer_debugger.py

@@ -20,6 +20,7 @@
 from absl import app
 from absl import flags
 from absl import logging
+from ai_edge_litert import interpreter as litert_interpreter
 import numpy as np
 import tensorflow as tf
 
@@ -194,7 +195,7 @@ def main(_) -> None:
       intrepreter_config=runtime.InterpreterConfig.kPreserveAllTensors,
   )
 
-  tflite_interpreter = tf.lite.Interpreter(
+  tflite_interpreter = litert_interpreter.Interpreter(
       model_path=_INPUT_TFLITE_FILE.value,
       experimental_preserve_all_tensors=True,
   )

third_party/python_requirements.in

@@ -34,3 +34,4 @@ mako
 pillow
 yapf
 protobuf
+ai-edge-litert

third_party/python_requirements.txt

@@ -11,6 +11,17 @@ absl-py==2.0.0 \
     #   keras
     #   tensorboard
     #   tensorflow
+ai-edge-litert==1.0.1 \
+    --hash=sha256:25a9b1577941498842bf77630722eda1163026c37abd57af66791a6955551b9d \
+    --hash=sha256:3bcec68832cdbdeb222aea43f31a70f9a3d55c17c1a0e4daf29c59109656c086 \
+    --hash=sha256:48f525fed0af268f4e904b0c9c8ab174e48422f9c2506b679c468a15e6bf64fc \
+    --hash=sha256:766c92511bb06851005454f164f10f4a2ed79f94877b11860c1e14158f948221 \
+    --hash=sha256:8c65c75a5b9f40a050416c62fd068e428848c96ec1873155976af884afac49e5 \
+    --hash=sha256:afa5eaca094f62e0b7c1981a94b1d03389cbc0e4a963eefb82d49053dd97bd73 \
+    --hash=sha256:e31a2aa9b031c42b1534744bbce78ef3fae7a7ac5a07008340554a5589d91003 \
+    --hash=sha256:ed8e7c3b058c2ba4233c1b3a0d59b9ca3197edf19d21be6710c8098472ab1e5b \
+    --hash=sha256:ef15c8cc1c029e047321541fd51932b745fba4ee21da64567c43d862b95687d9
+    # via -r third_party/python_requirements.in
 astunparse==1.6.3 \
     --hash=sha256:5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872 \
     --hash=sha256:c2652417f2c8b5bb325c885ae329bdf3f86424075c4fd1a128674bc6fba4b8e8
@@ -505,6 +516,7 @@ numpy==1.26.3 \
     --hash=sha256:f73497e8c38295aaa4741bdfa4fda1a5aedda5473074369eca10626835445511
     # via
     #   -r third_party/python_requirements.in
+    #   ai-edge-litert
     #   h5py
     #   keras
     #   ml-dtypes