@@ -64,15 +64,19 @@ This release contains contributions from many people at Google, as well as:
 ### Breaking Changes
 
 * `LiteRT`, a.k.a. `tf.lite`:
-  * C++ API:
-    * The public constants `tflite::Interpreter:kTensorsReservedCapacity`
-      and `tflite::Interpreter:kTensorsCapacityHeadroom` are now const
-      references, rather than `constexpr` compile-time constants.
-      (This is to enable better API compatibility for TFLite in Play services
-      while preserving the implementation flexibility to change the values of
-      these constants in the future.)
-    * Interpreter:
-      * `tf.lite.Interpreter` gives deprecation warning redirecting to its new location at `ai_edge_litert.interpreter`, as the API `tf.lite.Interpreter` will be deleted in TF 2.20. See the [migration guide](https://ai.google.dev/edge/litert/migration) for details.
+    * C++ API:
+      * The public constants `tflite::Interpreter:kTensorsReservedCapacity`
+        and `tflite::Interpreter:kTensorsCapacityHeadroom` are now const
+        references, rather than `constexpr` compile-time constants.
+        (This is to enable better API compatibility for TFLite in Play services
+        while preserving the implementation flexibility to change the values of
+        these constants in the future.)
+    * Python API:
+      * `tf.lite.Interpreter` gives deprecation warning redirecting to its new
+        location at `ai_edge_litert.interpreter`, as the API
+        `tf.lite.Interpreter` will be deleted in TF 2.20. See the
+        [migration guide](https://ai.google.dev/edge/litert/migration) for
+        details.
 
 ### Known Caveats
 
 
@@ -6,3 +6,6 @@ twine ~= 5.1.1
 build ~= 1.2.2
 # uv is faster than pip for installing Python packages.
 uv ~= 0.5.30
+
+# For running wheel verification script
+immutabledict ~= 4.2.1
@@ -221,6 +221,8 @@ void AllocateAndParseFlags() {
   build_ops_flags->tf_xla_check_cluster_output_numerics = false;
   build_ops_flags->tf_xla_disable_constant_folding = false;
   build_ops_flags->tf_xla_disable_full_embedding_pipelining = false;
+  build_ops_flags->tf_xla_disable_full_embedding_pipelining_with_summaries =
+      true;
   build_ops_flags->tf_xla_embedding_parallel_iterations = 0;
 
   mark_for_compilation_flags = new MarkForCompilationPassFlags;
@@ -312,6 +314,11 @@ void AllocateAndParseFlags() {
             &build_ops_flags->tf_xla_disable_full_embedding_pipelining,
             "If true then disables full embedding pipelining and instead use "
             "strict SparseCore / TensorCore sequencing."),
+       Flag("tf_xla_disable_full_embedding_pipelining_with_summaries",
+            &build_ops_flags
+                 ->tf_xla_disable_full_embedding_pipelining_with_summaries,
+            "If true then disables full embedding pipelining when summary ops "
+            "are detected."),
        Flag("tf_xla_embedding_parallel_iterations",
             &build_ops_flags->tf_xla_embedding_parallel_iterations,
             "If >0 then use this many parallel iterations in "
 
@@ -277,6 +277,10 @@ struct BuildXlaOpsPassFlags {
   // TensorCore sequencing will be used.
   bool tf_xla_disable_full_embedding_pipelining;
 
+  // Whether to enable automatical embedding pipelining when summary ops are
+  // detected in the graph.
+  bool tf_xla_disable_full_embedding_pipelining_with_summaries;
+
   // Force the WhileOps in embedding_pipelining and embedding_sequencing to use
   // this many parallel_iterations
   int tf_xla_embedding_parallel_iterations;
 
@@ -558,7 +558,8 @@ class Translator {
       OpOrArgNameMapper* op_or_arg_name_mapper,
       const std::map<std::string, std::string>& metadata,
       bool serialize_stablehlo_ops,
-      std::optional<size_t> custom_option_alignment);
+      std::optional<size_t> custom_option_alignment,
+      bool disable_buffer_deduping = false);
 
  private:
   enum class OpType : char { kTfliteBuiltin, kSelectTf, kCustomOp };
@@ -567,7 +568,8 @@ class Translator {
                       const std::unordered_set<std::string>& saved_model_tags,
                       OpOrArgNameMapper* op_or_arg_name_mapper,
                       const std::map<std::string, std::string>& metadata,
-                      std::optional<size_t> custom_option_alignment)
+                      std::optional<size_t> custom_option_alignment,
+                      bool disable_buffer_deduping)
       : module_(module),
         name_mapper_(*op_or_arg_name_mapper),
         builder_(kInitialBufferSize),
@@ -580,7 +582,8 @@ class Translator {
                             converter_flags.supported_backends().end()),
         serialize_debug_metadata_(converter_flags.serialize_debug_metadata()),
         use_buffer_offset_(converter_flags.use_buffer_offset()),
-        custom_option_alignment_(custom_option_alignment) {
+        custom_option_alignment_(custom_option_alignment),
+        disable_buffer_deduping_(disable_buffer_deduping) {
     // The first buffer must be empty according to the schema definition.
     empty_buffer_ = tflite::CreateBuffer(builder_);
     buffers_.push_back(empty_buffer_);
@@ -930,6 +933,8 @@ class Translator {
 
   std::optional<size_t> custom_option_alignment_ = std::nullopt;
 
+  bool disable_buffer_deduping_ = false;
+
   // Map from mlir constant attribute to the buffer index. This is used to
   // deduplicate the buffers in the flatbuffer.
   llvm::DenseMap<mlir::ElementsAttr, int> const_attribute_to_buffer_map_;
@@ -967,6 +972,7 @@ std::string Translator::UniqueName(mlir::Value val) {
 
 std::optional<BufferOffset<tflite::Buffer>> Translator::BuildBuffer(
     mlir::Value value, bool can_be_deduplicated, int& index) {
+  can_be_deduplicated = can_be_deduplicated && !disable_buffer_deduping_;
   auto inst = value.getDefiningOp();
   ElementsAttr attr;
   if (auto cst = dyn_cast<mlir::arith::ConstantOp>(inst)) {
@@ -3870,8 +3876,8 @@ std::optional<std::string> Translator::Translate(
     const std::unordered_set<std::string>& tags,
     OpOrArgNameMapper* op_or_arg_name_mapper,
     const std::map<std::string, std::string>& metadata,
-    bool serialize_stablehlo_ops,
-    std::optional<size_t> custom_option_alignment) {
+    bool serialize_stablehlo_ops, std::optional<size_t> custom_option_alignment,
+    bool disable_buffer_deduping) {
   OpOrArgLocNameMapper default_op_or_arg_name_mapper;
   if (!op_or_arg_name_mapper)
     op_or_arg_name_mapper = &default_op_or_arg_name_mapper;
@@ -3887,9 +3893,9 @@ std::optional<std::string> Translator::Translate(
     new_converter_flags.set_use_buffer_offset(true);
   }
 
-  auto translator = std::unique_ptr<Translator>(
-      new Translator(module, new_converter_flags, tags, op_or_arg_name_mapper,
-                     metadata, custom_option_alignment));
+  auto translator = std::unique_ptr<Translator>(new Translator(
+      module, new_converter_flags, tags, op_or_arg_name_mapper, metadata,
+      custom_option_alignment, disable_buffer_deduping));
   translator->convert_stablehlo_ = serialize_stablehlo_ops;
   auto ret = translator->TranslateInternal();
 
@@ -3900,9 +3906,9 @@ std::optional<std::string> Translator::Translate(
     ret = std::nullopt;
     auto new_converter_flags = converter_flags;
     new_converter_flags.set_use_buffer_offset(true);
-    translator = std::unique_ptr<Translator>(
-        new Translator(module, new_converter_flags, tags, op_or_arg_name_mapper,
-                       metadata, custom_option_alignment));
+    translator = std::unique_ptr<Translator>(new Translator(
+        module, new_converter_flags, tags, op_or_arg_name_mapper, metadata,
+        custom_option_alignment, disable_buffer_deduping));
     return translator->TranslateInternal();
   }
   return ret;
@@ -4495,7 +4501,7 @@ bool MlirToFlatBufferTranslateFunction(mlir::ModuleOp module,
   auto maybe_translated = Translator::Translate(
       module, options.converter_flags, options.saved_model_tags,
       options.op_or_arg_name_mapper, options.metadata, serialize_stablehlo_ops,
-      options.custom_option_alignment);
+      options.custom_option_alignment, options.disable_buffer_deduping);
   if (!maybe_translated) return false;
   *serialized_flatbuffer = std::move(*maybe_translated);
   return true;
 
@@ -48,6 +48,9 @@ struct FlatbufferExportOptions {
   // options. If specified, the value should be multiplier of 16 (default
   // alignment for TFL flatbuffer).
   std::optional<size_t> custom_option_alignment = std::nullopt;
+  // Whether to disable buffer-deduping which emits tensors with shared
+  // buffers.
+  bool disable_buffer_deduping = false;
 };
 
 // Translates the given MLIR `module` into a FlatBuffer and stores the
 
@@ -160,6 +160,12 @@ static opt<bool, true> serialize_debug_metadata_flag(
     llvm::cl::desc("Wether to serialize debug metadata or not"),
     llvm::cl::location(serialize_debug_metadata), llvm::cl::init(false));
 
+// NOLINTNEXTLINE
+static opt<bool> disable_buffer_deduping_flag(
+    "disable-buffer-deduping",
+    llvm::cl::desc("Whether to disable buffer deduping or not"),
+    llvm::cl::init(false));
+
 namespace mlir {
 namespace {
 static OwningOpRef<mlir::ModuleOp> FlatBufferFileToMlirTrans(
@@ -212,6 +218,7 @@ static LogicalResult MlirToFlatBufferFileTranslateFunction(
   options.op_or_arg_name_mapper = op_or_arg_name_mapper.get();
   options.converter_flags.set_serialize_debug_metadata(
       serialize_debug_metadata);
+  options.disable_buffer_deduping = disable_buffer_deduping_flag.getValue();
   if (!tflite::MlirToFlatBufferTranslateFunction(
           module, options, &serialized_flatbuffer, emit_stablehlo_ops))
     return mlir::failure();
 
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/lite/kernels/internal/runtime_shape.h"
 
-#include <algorithm>
 #include <cstdint>
 #include <functional>
 #include <initializer_list>
 
@@ -1,4 +1,5 @@
-// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s
+// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s --check-prefix=CHECK
+// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer -disable-buffer-deduping %s -o - | flatbuffer_to_string - | FileCheck %s --check-prefix=NO_DEDUPE
 
 module {
 func.func @add(%arg0: tensor<3x2xf32>) -> tensor<3x2xf32> attributes {tf.entry_function = {inputs = "serving_default_x", outputs = "outputs"}} {
@@ -90,4 +91,101 @@ func.func @sub(%arg0: tensor<3x2xf32>) -> tensor<3x2xf32> attributes {tf.entry_f
 // CHECK-NEXT:   }, {
 // CHECK-NEXT:     data: [ 49, 46, 54, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
 // CHECK-NEXT:   } ],
-// CHECK:      }
+// CHECK:      }
+
+// NO_DEDUPE: {
+// NO_DEDUPE:   version: 3,
+// NO_DEDUPE:   operator_codes: [ {
+// NO_DEDUPE:     version: 1
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:     deprecated_builtin_code: 41,
+// NO_DEDUPE:     version: 1,
+// NO_DEDUPE:     builtin_code: SUB
+// NO_DEDUPE:   } ],
+// NO_DEDUPE:   subgraphs: [ {
+// NO_DEDUPE:     tensors: [ {
+// NO_DEDUPE:       shape: [ 3, 2 ],
+// NO_DEDUPE:       buffer: 1,
+// NO_DEDUPE:       name: "serving_default_x",
+// NO_DEDUPE:       quantization: {
+// NO_DEDUPE:       },
+// NO_DEDUPE:       has_rank: true
+// NO_DEDUPE:     }, {
+// NO_DEDUPE:       shape: [ 3, 2 ],
+// NO_DEDUPE:       buffer: 2,
+// NO_DEDUPE:       name: "tfl.pseudo_const",
+// NO_DEDUPE:       quantization: {
+// NO_DEDUPE:       },
+// NO_DEDUPE:       has_rank: true
+// NO_DEDUPE:     }, {
+// NO_DEDUPE:       shape: [ 3, 2 ],
+// NO_DEDUPE:       buffer: 3,
+// NO_DEDUPE:       name: "outputs",
+// NO_DEDUPE:       quantization: {
+// NO_DEDUPE:       },
+// NO_DEDUPE:       has_rank: true
+// NO_DEDUPE:     } ],
+// NO_DEDUPE:     inputs: [ 0 ],
+// NO_DEDUPE:     outputs: [ 2 ],
+// NO_DEDUPE:     operators: [ {
+// NO_DEDUPE:       inputs: [ 1, 0 ],
+// NO_DEDUPE:       outputs: [ 2 ],
+// NO_DEDUPE:       builtin_options_type: AddOptions,
+// NO_DEDUPE:       builtin_options: {
+// NO_DEDUPE:       }
+// NO_DEDUPE:     } ],
+// NO_DEDUPE:     name: "add"
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:     tensors: [ {
+// NO_DEDUPE:       shape: [ 3, 2 ],
+// NO_DEDUPE:       buffer: 4,
+// NO_DEDUPE:       name: "serving_default_x",
+// NO_DEDUPE:       quantization: {
+// NO_DEDUPE:       },
+// NO_DEDUPE:       has_rank: true
+// NO_DEDUPE:     }, {
+// NO_DEDUPE:       shape: [ 3, 2 ],
+// NO_DEDUPE:       buffer: 5,
+// NO_DEDUPE:       name: "tfl.pseudo_const1",
+// NO_DEDUPE:       quantization: {
+// NO_DEDUPE:       },
+// NO_DEDUPE:       has_rank: true
+// NO_DEDUPE:     }, {
+// NO_DEDUPE:       shape: [ 3, 2 ],
+// NO_DEDUPE:       buffer: 6,
+// NO_DEDUPE:       name: "outputs",
+// NO_DEDUPE:       quantization: {
+// NO_DEDUPE:       },
+// NO_DEDUPE:       has_rank: true
+// NO_DEDUPE:     } ],
+// NO_DEDUPE:     inputs: [ 0 ],
+// NO_DEDUPE:     outputs: [ 2 ],
+// NO_DEDUPE:     operators: [ {
+// NO_DEDUPE:       opcode_index: 1,
+// NO_DEDUPE:       inputs: [ 1, 0 ],
+// NO_DEDUPE:       outputs: [ 2 ],
+// NO_DEDUPE:       builtin_options_type: SubOptions,
+// NO_DEDUPE:       builtin_options: {
+// NO_DEDUPE:       }
+// NO_DEDUPE:     } ],
+// NO_DEDUPE:     name: "sub"
+// NO_DEDUPE:   } ],
+// NO_DEDUPE:   description: "MLIR Converted.",
+// NO_DEDUPE:   buffers: [ {
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:     data: [ 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, 64 ]
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:     data: [ 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, 64 ]
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:   }, {
+// NO_DEDUPE:     data: [ 49, 46, 54, 46, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ]
+// NO_DEDUPE:   } ],
+// NO_DEDUPE:   metadata: [ {
+// NO_DEDUPE:     name: "min_runtime_version",
+// NO_DEDUPE:     buffer: 7
+// NO_DEDUPE:   } ],
+// NO_DEDUPE:   signature_defs: [  ]
+// NO_DEDUPE: }
@@ -181,3 +181,16 @@ func.func @FoldReshape(%arg0: tensor<4xi32>, %arg1: tensor<1x48x80x16x!quant.uni
   // CHECK{LITERAL}: "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x16x!quant.uniform<i8<-127:127>:f32, 0.022395913056501255>>, value = dense<[[[[12, -60, -51, -59, -62, 33, 53, 17, -31, 50, 27, 7, -19, -34, -14, -26], [47, -84, -32, -36, -102, -8, -8, 35, -33, 59, 95, 40, -25, -30, -55, 25]], [[4, -41, -61, 12, -23, 48, 40, 15, -39, 52, 81, -62, -24, 17, -7, -52], [40, -70, -45, 32, -43, 2, -30, 34, -35, 58, 77, -28, -30, 37, -47, -5]]]]> : tensor<1x2x2x16xi8>}> : () -> tensor<1x2x2x16x!quant.uniform<i8<-127:127>:f32, 0.022395913056501255>>
   // CHECK-NEXT: "tfl.transpose_conv"
 }
+
+// CHECK-LABEL: @FoldPerAxisReshape
+func.func @FoldPerAxisReshape() -> tensor<1x2x2x!quant.uniform<i8:f32:2, {0.007,0.004}>> {
+  %cst = arith.constant dense<[1, 2, 2]> : tensor<3xi32>
+  %0 = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i8:f32:1, {0.007,0.004}>>, value = dense<[[-127, 127], [-85, -80]]> : tensor<2x2xi8>}> : () -> tensor<2x2x!quant.uniform<i8:f32:1, {0.007,0.004}>>
+  %1 = "tfl.reshape"(%0, %cst) : (tensor<2x2x!quant.uniform<i8:f32:1, {0.007,0.004}>>, tensor<3xi32>) -> tensor<1x2x2x!quant.uniform<i8:f32:2, {0.007,0.004}>>
+  return %1 : tensor<1x2x2x!quant.uniform<i8:f32:2, {0.007,0.004}>>
+  
+
+// CHECK{LITERAL}:  %0 = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x!quant.uniform<i8:f32:2, {7.000000e-03,4.000000e-03}>>, value = dense<[[[-127, 127], [-85, -80]]]> : tensor<1x2x2xi8>}> : () -> tensor<1x2x2x!quant.uniform<i8:f32:2, {7.000000e-03,4.000000e-03}>>
+// CHECK-NOT: tfl.reshape
+// CHECK:  return %0 : tensor<1x2x2x!quant.uniform<i8:f32:2, {7.000000e-03,4.000000e-03}>>
+}