Merge branch 'main' into reduce-min

Author: mergify[bot]

.github/mergify.yml

@@ -1,3 +1,6 @@
+merge_queue:
+  max_parallel_checks: 1
+
 queue_rules:
   - name: default
     checks_timeout: 2 h

python/tflite_micro/BUILD

@@ -130,6 +130,7 @@ py_package(
     # in the tflm tree.
     packages = [
         "python.tflite_micro",
+        "tensorflow.lite.micro.compression",
         "tensorflow.lite.micro.tools.generate_test_for_model",
         "tensorflow.lite.python",
         "tensorflow.lite.tools.flatbuffer_utils",
@@ -138,6 +139,7 @@ py_package(
         ":postinstall_check",
         ":runtime",
         ":version",
+        "//tensorflow/lite/micro/compression",
     ],
 )
 
@@ -223,8 +225,10 @@ py_wheel(
         ":local": "py3",
     }),
     requires = [
+        "bitarray",
         "flatbuffers",
         "numpy",
+        "pyyaml",
         "tensorflow",
     ],
     stamp = 1,  # 1 == always stamp

python/tflite_micro/__init__.py

@@ -24,5 +24,8 @@
 # Unambiguously identify the source used to build the package.
 from tflite_micro.python.tflite_micro._version import __version__
 
-# Ordered after `runtime` to avoid a circular dependency
+# Provide a convenient alias for the compression module
+from tflite_micro.tensorflow.lite.micro import compression
+
+# Ordered after `runtime` and `compression` to avoid circular dependencies
 from tflite_micro.python.tflite_micro import postinstall_check

python/tflite_micro/postinstall_check.py

@@ -19,13 +19,17 @@
 # in the Python installation environment rather than to locations in the tflm
 # source tree.
 from tflite_micro import runtime
+from tflite_micro import compression
 
 import numpy as np
 import pkg_resources
 import sys
+import tempfile
+import os
 
 
-def passed():
+def runtime_test():
+  """Test the runtime interpreter functionality."""
   # Create an interpreter with a sine model
   model = pkg_resources.resource_filename(__name__, "sine_float.tflite")
   interpreter = runtime.Interpreter.from_file(model)
@@ -49,5 +53,27 @@ def infer(x):
   return np.allclose(outputs, goldens, atol=0.05)
 
 
+def compression_test():
+  """Test that the compression module is available and functional."""
+
+  # Test that compress function is available
+  # We don't actually compress here as it requires a properly structured model
+  # with compressible tensors, but we verify the function is importable
+  assert callable(compression.compress)
+
+  # Test availability of the SpecBuilder
+  _ = (compression.SpecBuilder().add_tensor(
+      subgraph=0, tensor=0).with_lut(index_bitwidth=4).build())
+
+  return True
+
+
+def passed():
+  """Run all postinstall checks."""
+  runtime_passed = runtime_test()
+  compression_passed = compression_test()
+  return runtime_passed and compression_passed
+
+
 if __name__ == "__main__":
   sys.exit(0 if passed() else 1)

python/tflite_micro/pypi_build.sh

@@ -101,13 +101,15 @@ docker run \
     }
 
     # Build the wheel via bazel, using the Python compatibility tag matching the
-    # build environment.
+    # build environment. Enable compression support for the official package.
     call_bazel build //python/tflite_micro:whl.dist \
-        --//python/tflite_micro:compatibility_tag=\$PY_COMPATIBILITY
+        --//python/tflite_micro:compatibility_tag=\$PY_COMPATIBILITY \
+        --//:with_compression=true
 
     # Test, in the container environment.
     call_bazel test //python/tflite_micro:whl_test \
-            --//python/tflite_micro:compatibility_tag=\$PY_COMPATIBILITY
+            --//python/tflite_micro:compatibility_tag=\$PY_COMPATIBILITY \
+            --//:with_compression=true
 EOF
 
 # Make the output directory tree writable so it can be removed easily by the

python/tflite_micro/python_ops_resolver.cc

@@ -105,6 +105,7 @@ PythonOpsResolver::PythonOpsResolver() {
   AddReshape();
   AddResizeBilinear();
   AddResizeNearestNeighbor();
+  AddReverseV2();
   AddRfft();
   AddRound();
   AddRsqrt();

tensorflow/lite/core/c/common.cc

@@ -104,6 +104,41 @@ void TfLiteVarArrayFree(T* a) {
 
 #ifndef TF_LITE_STATIC_MEMORY
 
+TfLiteSparsity TfLiteSparsityClone(const TfLiteSparsity& src) {
+  TfLiteSparsity dst = src;
+  dst.traversal_order = TfLiteIntArrayCopy(src.traversal_order);
+  dst.block_map = TfLiteIntArrayCopy(src.block_map);
+  if (src.dim_metadata) {
+    dst.dim_metadata = reinterpret_cast<TfLiteDimensionMetadata*>(
+        calloc(1, sizeof(TfLiteDimensionMetadata) * src.dim_metadata_size));
+    for (int i = 0; i < src.dim_metadata_size; ++i) {
+      dst.dim_metadata[i] = src.dim_metadata[i];
+      dst.dim_metadata[i].array_segments =
+          TfLiteIntArrayCopy(src.dim_metadata[i].array_segments);
+      dst.dim_metadata[i].array_indices =
+          TfLiteIntArrayCopy(src.dim_metadata[i].array_indices);
+    }
+  }
+  return dst;
+}
+
+// Clones the source sparsity to a newly allocated object.
+TfLiteSparsity* TfLiteSparsityClone(const TfLiteSparsity* const src) {
+  if (!src) {
+    return nullptr;
+  }
+  TfLiteSparsity* dst =
+      reinterpret_cast<TfLiteSparsity*>(calloc(1, sizeof(TfLiteSparsity)));
+  *dst = TfLiteSparsityClone(*src);
+  return dst;
+}
+
+#endif  // TF_LITE_STATIC_MEMORY
+
+}  // namespace
+
+#ifndef TF_LITE_STATIC_MEMORY
+
 TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) {
   TfLiteQuantization dst;
   dst.type = src.type;
@@ -136,39 +171,8 @@ TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src) {
   return dst;
 }
 
-TfLiteSparsity TfLiteSparsityClone(const TfLiteSparsity& src) {
-  TfLiteSparsity dst = src;
-  dst.traversal_order = TfLiteIntArrayCopy(src.traversal_order);
-  dst.block_map = TfLiteIntArrayCopy(src.block_map);
-  if (src.dim_metadata) {
-    dst.dim_metadata = reinterpret_cast<TfLiteDimensionMetadata*>(
-        calloc(1, sizeof(TfLiteDimensionMetadata) * src.dim_metadata_size));
-    for (int i = 0; i < src.dim_metadata_size; ++i) {
-      dst.dim_metadata[i] = src.dim_metadata[i];
-      dst.dim_metadata[i].array_segments =
-          TfLiteIntArrayCopy(src.dim_metadata[i].array_segments);
-      dst.dim_metadata[i].array_indices =
-          TfLiteIntArrayCopy(src.dim_metadata[i].array_indices);
-    }
-  }
-  return dst;
-}
-
-// Clones the source sparsity to a newly allocated object.
-TfLiteSparsity* TfLiteSparsityClone(const TfLiteSparsity* const src) {
-  if (!src) {
-    return nullptr;
-  }
-  TfLiteSparsity* dst =
-      reinterpret_cast<TfLiteSparsity*>(calloc(1, sizeof(TfLiteSparsity)));
-  *dst = TfLiteSparsityClone(*src);
-  return dst;
-}
-
 #endif  // TF_LITE_STATIC_MEMORY
 
-}  // namespace
-
 extern "C" {
 
 size_t TfLiteIntArrayGetSizeInBytes(int size) {
@@ -247,6 +251,11 @@ void TfLiteQuantizationFree(TfLiteQuantization* quantization) {
     }
     free(q_params);
   }
+  if (quantization->type == kTfLiteBlockwiseQuantization) {
+    TfLiteBlockwiseQuantization* q_params =
+        reinterpret_cast<TfLiteBlockwiseQuantization*>(quantization->params);
+    free(q_params);
+  }
   quantization->params = nullptr;
   quantization->type = kTfLiteNoQuantization;
 }
@@ -640,4 +649,8 @@ TfLiteRunStep TfLiteTensorGetShapeKnownStep(const TfLiteTensor* t) {
   return kTfLiteRunStepUnknown;
 }
 
+// Returns a sentinel value to be used as the user_data field of a TfLiteNode
+// when the kernel initialization fails.
+void* TfLiteKernelInitFailed() { return reinterpret_cast<void*>(-1); }
+
 }  // extern "C"

tensorflow/lite/core/c/common.h

@@ -788,6 +788,7 @@ TfLiteStatus TfLiteTensorRealloc(size_t num_bytes, TfLiteTensor* tensor);
 /// If all dimensions are known, this is the same as `t->dims`.
 /// (`dims_signature` is NULL or empty if all dimensions are known.)
 const TfLiteIntArray* TfLiteTensorGetDimsSignature(const TfLiteTensor* t);
+
 #endif  // TF_LITE_STATIC_MEMORY
 
 /// WARNING: This is an experimental interface that is subject to change.
@@ -1160,6 +1161,11 @@ typedef struct TfLiteRegistration {
   /// NOTE: if the data is already in the desired format, simply implement this
   /// function to return `nullptr` and implement the free function to be a
   /// no-op.
+  ///
+  /// NOTE: For a Delegate kernel, returns `TfLiteKernelInitFailed()` if it
+  /// fails on the initialization. This eventually causes user's API call to
+  /// InterpreterBuilder::operator() or Interpreter::ModifyGraphWithDelegate()
+  /// to return an error.
   void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
 
   /// The pointer `buffer` is the data previously returned by an init
@@ -1498,6 +1504,10 @@ TfLiteRunStep TfLiteTensorGetDataKnownStep(const TfLiteTensor* t);
 /// operations.
 TfLiteRunStep TfLiteTensorGetShapeKnownStep(const TfLiteTensor* t);
 
+/// Returns a sentinel value to be used as the user_data field of a TfLiteNode
+/// when the kernel initialization fails.
+void* TfLiteKernelInitFailed();
+
 /** @} */
 // Ends `\addtogroup`, it's important for the doc generator that this doesn't
 // include the CC code below.
@@ -1633,5 +1643,8 @@ TfLiteStatus TfLiteTensorVariantRealloc(TfLiteTensor* t,
   return kTfLiteOk;
 }
 
+// Returns a copy of the quantization parameters of the tensor.
+TfLiteQuantization TfLiteQuantizationClone(const TfLiteQuantization& src);
+
 #endif  // __cplusplus
 #endif  // TENSORFLOW_LITE_CORE_C_COMMON_H_

tensorflow/lite/kernels/internal/portable_tensor_utils.h

@@ -170,6 +170,7 @@ inline void BatchQuantizeFloats(const float* float_data_ptr, int n_batch,
       tensor_utils::SymmetricQuantizeFloats(
           float_data_ptr + offset, n_data, quantized_data_ptr + offset,
           &unused_min, &unused_max, &scaling_factors[b]);
+      if (zero_points) zero_points[b] = 0;
     }
   }
 }

tensorflow/lite/kernels/internal/reference/batch_matmul.h

@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/common.h"
 #include "tensorflow/lite/kernels/internal/compatibility.h"
 #include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
 #include "tensorflow/lite/kernels/internal/types.h"
 
 namespace tflite {