chore: import upstream snapshot with attribution

tensorflow/lite/tools/serialization/BUILD

@@ -0,0 +1,112 @@
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    features = ["-parse_headers"],
+    licenses = ["notice"],
+)
+
+cc_binary(
+    name = "option_writer_generator",
+    srcs = ["option_writer_generator.cc"],
+    deps = [
+        "//tensorflow/lite/schema:schema_fbs_with_reflection",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "writer_lib_enum",
+    hdrs = ["enum_mapping.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
+        "//tensorflow/lite:builtin_op_data",
+    ],
+)
+
+cc_library(
+    name = "writer_lib",
+    srcs = [
+        "writer_lib.cc",
+    ],
+    hdrs = [
+        "writer_lib.h",
+    ],
+    data = [
+        ":option_writer_gen",
+    ],
+    textual_hdrs = ["option_writer_generated.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/lite/schema:schema_conversion_utils",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
+        "//tensorflow/compiler/mlir/lite/tools/versioning",
+        "//tensorflow/lite:builtin_op_data",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:schema_fbs_version",
+        "//tensorflow/lite/core:framework_stable",
+        "//tensorflow/lite/core:model_builder",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/schema:schema_conversion_utils",
+        "//tensorflow/lite/tools/serialization:writer_lib_enum",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@flatbuffers//:runtime_cc",
+    ],
+)
+
+cc_binary(
+    name = "writer",
+    srcs = ["writer.cc"],
+    deps = [
+        ":writer_lib",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+    ],
+)
+
+cc_binary(
+    name = "writer_test",
+    srcs = ["writer_test.cc"],
+    deps = [
+        ":writer_lib",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+    ],
+)
+
+cc_test(
+    name = "writer_lib_test",
+    size = "small",
+    srcs = ["writer_lib_test.cc"],
+    deps = [
+        ":writer_lib",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
+        "//tensorflow/lite:builtin_ops",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/kernels:subgraph_test_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/log:check",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:logging",
+    ],
+)
+
+genrule(
+    name = "option_writer_gen",
+    outs = ["option_writer_generated.h"],
+    cmd = "$(location :option_writer_generator) $(@)",
+    tools = [":option_writer_generator"],
+)

tensorflow/lite/tools/serialization/README.md

@@ -0,0 +1,63 @@
+# TFLite Serialization Tool
+
+**NOTE:** This tool is intended for advanced users only, and should be used with
+care.
+
+The (C++) serialization library generates and writes a TFLite flatbuffer given
+an `Interpreter` or `Subgraph`. Example use-cases include authoring models with
+the `Interpreter` API, or updating models on-device (by modifying `tensor.data`
+for relevant tensors).
+
+## Serialization
+
+### Writing flatbuffer to file
+
+To write a TFLite model from an `Interpreter` (see `lite/interpreter.h`):
+`std::unique_ptr<tflite::Interpreter> interpreter; // ...build/modify
+interpreter... tflite::ModelWriter writer(interpreter.get()); std::string
+filename = "/tmp/model.tflite"; writer.Write(filename);`
+
+Note that the above API does not support custom I/O tensors or custom ops yet.
+However, it does support model with Control Flow.
+
+To generate/write a flatbuffer for a particular `Subgraph` (see
+`lite/core/subgraph.h`) you can use `SubgraphWriter`.
+
+```
+std::unique_ptr<tflite::Interpreter> interpreter;
+// ...build/modify interpreter...
+// The number of subgraphs can be obtained by:
+// const int num_subgraphs = interpreter_->subgraphs_size();
+// Note that 0 <= subgraph_index < num_subgraphs
+tflite::SubgraphWriter writer(&interpreter->subgraph(subgraph_index));
+std::string filename = "/tmp/model.tflite";
+writer.Write(filename);
+```
+
+`SubgraphWriter` supports custom ops and/or custom I/O tensors.
+
+### Generating flatbuffer in-memory
+
+Both `ModelWriter` and `SubgraphWriter` support a `GetBuffer` method to return
+the generated flatbuffer in-memory:
+
+```
+std::unique_ptr<uint8_t[]> output_buffer;
+size_t output_buffer_size;
+tflite::ModelWriter writer(interpreter.get());
+writer.GetBuffer(&output_buffer, &output_buffer_size);
+```
+
+## De-serialization
+
+The flatbuffers written as above can be de-serialized just like any other TFLite
+model, for eg:
+
+```
+std::unique_ptr<FlatBufferModel> model =
+    FlatBufferModel::BuildFromFile(filename);
+tflite::ops::builtin::BuiltinOpResolver resolver;
+InterpreterBuilder builder(*model, resolver);
+std::unique_ptr<Interpreter> new_interpreter;
+builder(&new_interpreter);
+```

tensorflow/lite/tools/serialization/enum_mapping.h

@@ -0,0 +1,170 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_TOOLS_SERIALIZATION_ENUM_MAPPING_H_
+#define TENSORFLOW_LITE_TOOLS_SERIALIZATION_ENUM_MAPPING_H_
+
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
+#include "tensorflow/lite/builtin_op_data.h"
+
+// TODO(aselle): Ideally extract this from the schema.
+
+namespace tflite {
+
+inline ActivationFunctionType TfLiteActivationToSchemaActivation(
+    TfLiteFusedActivation act) {
+  switch (act) {
+    case kTfLiteActNone:
+      return ActivationFunctionType_NONE;
+    case kTfLiteActRelu:
+      return ActivationFunctionType_RELU;
+    case kTfLiteActReluN1To1:
+      return ActivationFunctionType_RELU_N1_TO_1;
+    case kTfLiteActRelu6:
+      return ActivationFunctionType_RELU6;
+    case kTfLiteActTanh:
+      return ActivationFunctionType_TANH;
+    case kTfLiteActSignBit:
+      return ActivationFunctionType_SIGN_BIT;
+    case kTfLiteActSigmoid:
+      return ActivationFunctionType_NONE;  // TODO(aselle): Add to schema
+  }
+  return ActivationFunctionType_NONE;
+}
+
+inline Padding TfLitePaddingToSchemaPadding(TfLitePadding padding) {
+  switch (padding) {
+    case kTfLitePaddingUnknown:
+      return Padding_SAME;  // TODO(aselle): Consider an error.
+    case kTfLitePaddingSame:
+      return Padding_SAME;
+    case kTfLitePaddingValid:
+      return Padding_VALID;
+  }
+  return Padding_SAME;  // TODO(aselle): Consider an error.
+}
+
+inline TensorType TfLiteTypeToSchemaType(TfLiteType type) {
+  switch (type) {
+    // case kTfLiteNoType: return TensorType_NONE;
+    case kTfLiteNoType:
+      return TensorType_FLOAT32;  // TODO(aselle): Consider an error.
+    case kTfLiteFloat32:
+      return TensorType_FLOAT32;
+    case kTfLiteFloat16:
+      return TensorType_FLOAT16;
+    case kTfLiteBFloat16:
+      return TensorType_BFLOAT16;
+    case kTfLiteFloat64:
+      return TensorType_FLOAT64;
+    case kTfLiteInt32:
+      return TensorType_INT32;
+    case kTfLiteUInt32:
+      return TensorType_UINT32;
+    case kTfLiteUInt4:
+      return TensorType_UINT4;
+    case kTfLiteFloat8E4M3FN:
+      return TensorType_FLOAT8_E4M3FN;
+    case kTfLiteFloat8E5M2:
+      return TensorType_FLOAT8_E5M2;
+    case kTfLiteInt4:
+      return TensorType_INT4;
+    case kTfLiteInt2:
+      return TensorType_INT2;
+    case kTfLiteUInt8:
+      return TensorType_UINT8;
+    case kTfLiteInt8:
+      return TensorType_INT8;
+    case kTfLiteInt64:
+      return TensorType_INT64;
+    case kTfLiteUInt64:
+      return TensorType_UINT64;
+    case kTfLiteString:
+      return TensorType_STRING;
+    case kTfLiteBool:
+      return TensorType_BOOL;
+    case kTfLiteUInt16:
+      return TensorType_UINT16;
+    case kTfLiteInt16:
+      return TensorType_INT16;
+    case kTfLiteComplex64:
+      return TensorType_COMPLEX64;
+    case kTfLiteComplex128:
+      return TensorType_COMPLEX128;
+    case kTfLiteResource:
+      return TensorType_RESOURCE;
+    case kTfLiteVariant:
+      return TensorType_VARIANT;
+  }
+  // TODO(aselle): consider an error
+}
+
+inline FullyConnectedOptionsWeightsFormat
+FullyConnectedOptionsWeightsFormatToSchema(
+    TfLiteFullyConnectedWeightsFormat format) {
+  switch (format) {
+    case kTfLiteFullyConnectedWeightsFormatDefault:
+      return FullyConnectedOptionsWeightsFormat_DEFAULT;
+    case kTfLiteFullyConnectedWeightsFormatShuffled4x16Int8:
+      return FullyConnectedOptionsWeightsFormat_SHUFFLED4x16INT8;
+  }
+}
+
+inline LSTMKernelType LSTMKernelTypeToSchema(TfLiteLSTMKernelType type) {
+  switch (type) {
+    case kTfLiteLSTMFullKernel:
+      return LSTMKernelType_FULL;
+    case kTfLiteLSTMBasicKernel:
+      return LSTMKernelType_BASIC;
+  }
+}
+
+inline LSHProjectionType LSHProjectionTypeToSchema(
+    TfLiteLSHProjectionType type) {
+  switch (type) {
+    case kTfLiteLshProjectionUnknown:
+      return LSHProjectionType_UNKNOWN;
+    case kTfLiteLshProjectionSparse:
+      return LSHProjectionType_SPARSE;
+    case kTfLiteLshProjectionDense:
+      return LSHProjectionType_DENSE;
+  }
+}
+
+inline MirrorPadMode MirrorPaddingModeToSchema(TfLiteMirrorPaddingMode mode) {
+  switch (mode) {
+    case kTfLiteMirrorPaddingUnknown:
+      return MirrorPadMode_REFLECT;  // TODO(aselle): consider an error
+    case kTfLiteMirrorPaddingReflect:
+      return MirrorPadMode_REFLECT;
+    case kTfLiteMirrorPaddingSymmetric:
+      return MirrorPadMode_SYMMETRIC;
+  }
+}
+
+inline CombinerType CombinerTypeToSchema(TfLiteCombinerType type) {
+  switch (type) {
+    case kTfLiteCombinerTypeSum:
+      return CombinerType_SUM;
+    case kTfLiteCombinerTypeMean:
+      return CombinerType_MEAN;
+    case kTfLiteCombinerTypeSqrtn:
+      return CombinerType_SQRTN;
+  }
+}
+
+// int
+
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_TOOLS_SERIALIZATION_ENUM_MAPPING_H_

tensorflow/lite/tools/serialization/option_writer_generator.cc

@@ -0,0 +1,586 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+#include <ctype.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/schema/reflection/schema_generated.h"
+
+namespace tflite {
+namespace {
+// This is generated by grepping
+//  cat third_party/tensorflow/lite/core/c/builtin_op_data.h | grep "^} TfLite"
+//  | sed 's/^} \(TfLite.*\)Params;/\1Params/g' | grep -v "^}" | sed
+//  's/\(.*\)/"\1",/g' | sort
+static const char* param_structs[] = {"TfLiteAddParams",
+                                      "TfLiteArgMaxParams",
+                                      "TfLiteArgMinParams",
+                                      "TfLiteBatchMatMulParams",
+                                      "TfLiteBatchToSpaceNDParams",
+                                      "TfLiteBidirectionalSequenceLSTMParams",
+                                      "TfLiteBidirectionalSequenceRNNParams",
+                                      "TfLiteBucketizeParams",
+                                      "TfLiteCastParams",
+                                      "TfLiteConcatenationParams",
+                                      "TfLiteConvParams",
+                                      "TfLiteDepthwiseConvParams",
+                                      "TfLiteDivParams",
+                                      "TfLiteDynamicUpdateSliceParams",
+                                      "TfLiteEmbeddingLookupSparseParams",
+                                      "TfLiteFakeQuantParams",
+                                      "TfLiteFullyConnectedParams",
+                                      "TfLiteGatherParams",
+                                      "TfLiteGeluParams",
+                                      "TfLiteIfParams",
+                                      "TfLiteL2NormParams",
+                                      "TfLiteLeakyReluParams",
+                                      "TfLiteLocalResponseNormParams",
+                                      "TfLiteLSHProjectionParams",
+                                      "TfLiteLSTMParams",
+                                      "TfLiteMirrorPaddingParams",
+                                      "TfLiteMulParams",
+                                      "TfLiteOneHotParams",
+                                      "TfLitePackParams",
+                                      "TfLitePadParams",
+                                      "TfLitePadV2Params",
+                                      "TfLitePoolParams",
+                                      "TfLiteRandomParams",
+                                      "TfLiteReducerParams",
+                                      "TfLiteReshapeParams",
+                                      "TfLiteResizeBilinearParams",
+                                      "TfLiteResizeNearestNeighborParams",
+                                      "TfLiteRNNParams",
+                                      "TfLiteSequenceRNNParams",
+                                      "TfLiteShapeParams",
+                                      "TfLiteSkipGramParams",
+                                      "TfLiteSoftmaxParams",
+                                      "TfLiteSpaceToBatchNDParams",
+                                      "TfLiteSpaceToDepthParams",
+                                      "TfLiteDepthToSpaceParams",
+                                      "TfLiteSparseToDenseParams",
+                                      "TfLiteSplitParams",
+                                      "TfLiteSplitVParams",
+                                      "TfLiteSqueezeParams",
+                                      "TfLiteStablehloScatterParams",
+                                      "TfLiteStridedSliceParams",
+                                      "TfLiteSubParams",
+                                      "TfLiteSVDFParams",
+                                      "TfLiteTransposeConvParams",
+                                      "TfLiteTransposeParams",
+                                      "TfLiteUnidirectionalSequenceLSTMParams",
+                                      "TfLiteUniqueParams",
+                                      "TfLiteUnpackParams",
+                                      "TfLiteReverseSequenceParams",
+                                      "TfLiteWhileParams",
+                                      "TfLiteCumsumParams",
+                                      "TfLiteCallOnceParams",
+                                      "TfLiteConv3DParams",
+                                      "TfLiteHashtableParams",
+                                      "TfLiteHashtableFindParams",
+                                      "TfLiteHashtableImportParams",
+                                      "TfLiteHashtableSizeParams",
+                                      "TfLiteConv3DTransposeParams",
+                                      "TfLiteVarHandleParams",
+                                      "TfLiteUnsortedSegmentSumParams",
+                                      "TfLiteUnsortedSegmentMinParams",
+                                      "TfLiteBitwiseXorParams",
+                                      "TfLiteRightShiftParams",
+                                      "TfLiteReduceWindowParams",
+                                      nullptr};
+}  // namespace
+
+// Get rid of all underscores and make everything lower case to make name
+// matching work for stuff like 3D vs 3d or RNN vs Rnn.
+std::string ToCollapsed(const std::string& in) {
+  const char* s = in.c_str();
+  bool first = true;
+  std::string out;
+  while (*s != '\0') {
+    if (*s == '_') {
+      first = true;
+    } else if (first) {
+      out.push_back(tolower(*s));
+      first = false;
+    } else {
+      out.push_back(tolower(*s));
+    }
+    s++;
+  }
+  return out;
+}
+
+// A collection of information about builtin ops.
+class OpOptionData {
+ public:
+  OpOptionData() {
+    BuildOpList();
+    BuildOptionToTypeFunctionMap();
+    BuildOpToOptionMap();
+  }
+
+  // A list of builtin operations
+  const std::vector<std::string>& ops() const { return ops_; }
+  // Maps from operation name to option name (i.e. 'ADD' to 'AddOptions')
+  const std::unordered_map<std::string, std::string>& op_to_option() {
+    return op_to_option_;
+  }
+  // Maps from option to C struct i.e. 'AddOptions' -> 'TfLiteAddOptions'
+  const std::unordered_map<std::string, std::string>& option_to_struct() {
+    return option_to_struct_;
+  }
+  // Maps from option to a flatbuffer type function that describes that option.
+  const std::unordered_map<std::string, flatbuffers::TypeFunction>&
+  option_to_type_function() {
+    return option_to_type_function_;
+  }
+
+ private:
+  void BuildOpList() {
+    for (const char* const* curr = EnumNamesBuiltinOperator(); *curr != nullptr;
+         ++curr) {
+      if (strlen(*curr) != 0) ops_.push_back(*curr);
+    }
+  }
+
+  void BuildOptionToTypeFunctionMap() {
+    auto d = tflite::BuiltinOptionsTypeTable();
+    for (int i = 0; i < d->num_elems; i++) {
+      flatbuffers::TypeCode code = d->type_codes[i];
+      if (code.sequence_ref != -1) {
+        option_to_type_function_.insert(
+            std::make_pair(d->names[i], d->type_refs[code.sequence_ref]));
+      }
+    }
+  }
+
+  void BuildOpToOptionMap() {
+    // Manually specified mappings between ops and options
+    op_to_option_["REDUCE_MAX"] = "ReducerOptions";
+    op_to_option_["REDUCE_MIN"] = "ReducerOptions";
+    op_to_option_["REDUCE_ANY"] = "ReducerOptions";
+    op_to_option_["REDUCE_ALL"] = "ReducerOptions";
+    op_to_option_["SUM"] = "ReducerOptions";
+    op_to_option_["REDUCE_MAX"] = "ReducerOptions";
+    op_to_option_["REDUCE_PROD"] = "ReducerOptions";
+    op_to_option_["MEAN"] = "ReducerOptions";
+    op_to_option_["L2_POOL_2D"] = "Pool2DOptions";
+    op_to_option_["AVERAGE_POOL_2D"] = "Pool2DOptions";
+    op_to_option_["MAX_POOL_2D"] = "Pool2DOptions";
+    op_to_option_["L2_NORMALIZATION"] = "L2NormOptions";
+    op_to_option_["UNIDIRECTIONAL_SEQUENCE_RNN"] = "SequenceRNNOptions";
+    op_to_option_["MAXIMUM"] = "MaximumMinimumOptions";
+    op_to_option_["MINIMUM"] = "MaximumMinimumOptions";
+    op_to_option_["CONV_3D_TRANSPOSE"] = "Conv3DOptions";
+    op_to_option_["RANDOM_STANDARD_NORMAL"] = "RandomOptions";
+    op_to_option_["RANDOM_UNIFORM"] = "RandomOptions";
+    op_to_option_["MULTINOMIAL"] = "RandomOptions";
+
+    // These operators are not real ones.
+    op_to_option_["CUSTOM"] = "";    // TODO(aselle): maybe something else.
+    op_to_option_["DELEGATE"] = "";  // TODO(aselle): maybe something else.
+    op_to_option_["PLACEHOLDER_FOR_GREATER_OP_CODES"] = "";
+
+    // Manually specified mappings between ops to "none" options -- these are
+    // ops without a corresponding Options message in schema as yet. If these
+    // options do get assigned an Options message in future, they need to be
+    // updated here as well.
+    op_to_option_["EMBEDDING_LOOKUP"] = "";
+    op_to_option_["FLOOR"] = "";
+    op_to_option_["CEIL"] = "";
+    op_to_option_["HASHTABLE_LOOKUP"] = "";
+    op_to_option_["LOGISTIC"] = "";
+    op_to_option_["RELU"] = "";
+    op_to_option_["RELU_N1_TO_1"] = "";
+    op_to_option_["RELU_0_TO_1"] = "";
+    op_to_option_["RELU6"] = "";
+    op_to_option_["ROUND"] = "";
+    op_to_option_["TANH"] = "";
+    op_to_option_["PRELU"] = "";
+    op_to_option_["SIN"] = "";
+    op_to_option_["LOG"] = "";
+    op_to_option_["SQRT"] = "";
+    op_to_option_["RSQRT"] = "";
+    op_to_option_["ELU"] = "";
+    op_to_option_["REVERSE_SEQUENCE"] = "";
+    op_to_option_["REAL"] = "";
+    op_to_option_["IMAG"] = "";
+    op_to_option_["COMPLEX_ABS"] = "";
+    op_to_option_["BROADCAST_ARGS"] = "";
+    op_to_option_["GELU"] = "";
+    op_to_option_["DYNAMIC_UPDATE_SLICE"] = "";
+    op_to_option_["BITCAST"] = "";
+    op_to_option_["BITWISE_XOR"] = "";
+    op_to_option_["RIGHT_SHIFT"] = "";
+    op_to_option_["STABLEHLO_LOGISTIC"] = "";
+    op_to_option_["STABLEHLO_ADD"] = "";
+    op_to_option_["STABLEHLO_DIVIDE"] = "";
+    op_to_option_["STABLEHLO_MULTIPLY"] = "";
+    op_to_option_["STABLEHLO_MAXIMUM"] = "";
+    op_to_option_["STABLEHLO_MINIMUM"] = "";
+    op_to_option_["STABLEHLO_RESHAPE"] = "";
+    op_to_option_["STABLEHLO_CLAMP"] = "";
+    op_to_option_["STABLEHLO_ABS"] = "";
+    op_to_option_["STABLEHLO_AND"] = "";
+    op_to_option_["STABLEHLO_COSINE"] = "";
+    op_to_option_["STABLEHLO_EXPONENTIAL"] = "";
+    op_to_option_["STABLEHLO_FLOOR"] = "";
+    op_to_option_["STABLEHLO_LOG"] = "";
+    op_to_option_["STABLEHLO_OR"] = "";
+    op_to_option_["STABLEHLO_NEGATE"] = "";
+    op_to_option_["STABLEHLO_POWER"] = "";
+    op_to_option_["STABLEHLO_REMAINDER"] = "";
+    op_to_option_["STABLEHLO_RSQRT"] = "";
+    op_to_option_["STABLEHLO_SELECT"] = "";
+    op_to_option_["STABLEHLO_SUBTRACT"] = "";
+    op_to_option_["STABLEHLO_TANH"] = "";
+    op_to_option_["STABLEHLO_CONVERT"] = "";
+    op_to_option_["STABLEHLO_DYNAMIC_UPDATE_SLICE"] = "";
+    op_to_option_["STABLEHLO_SHIFT_LEFT"] = "";
+    op_to_option_["STABLEHLO_CBRT"] = "";
+
+    // TODO(aselle): These are undesirable hacks. Consider changing C structs
+    option_to_struct_["Pool2DOptions"] = "TfLitePoolParams";
+    option_to_struct_["Conv2DOptions"] = "TfLiteConvParams";
+    option_to_struct_["DepthwiseConv2DOptions"] = "TfLiteDepthwiseConvParams";
+    option_to_struct_["LocalResponseNormalizationOptions"] =
+        "TfLiteLocalResponseNormParams";
+    option_to_struct_["MirrorPadOptions"] = "TfLiteMirrorPaddingParams";
+    // Now for every op, try to find an option.
+    bool fatal = false;
+    for (const auto& op_name : ops_) {
+      auto d = tflite::BuiltinOptionsTypeTable();
+      std::string collapsed_option_name_guess =
+          ToCollapsed(op_name) + "options";
+      // O(n^2) but not that big of n.
+      for (int i = 0; i < d->num_elems; i++) {
+        std::string option_name = d->names[i];
+        std::string collapsed_option_name = ToCollapsed(option_name);
+        if (collapsed_option_name_guess == collapsed_option_name) {
+          op_to_option_.insert(std::make_pair(op_name, option_name));
+          break;
+        }
+      }
+      // add second union into the check
+      auto d2 = tflite::BuiltinOptions2TypeTable();
+      for (int i = 0; i < d2->num_elems; i++) {
+        std::string option_name = d2->names[i];
+        std::string collapsed_option_name = ToCollapsed(option_name);
+        if (collapsed_option_name_guess == collapsed_option_name) {
+          op_to_option_.insert(std::make_pair(op_name, option_name));
+          break;
+        }
+      }
+
+      auto it = op_to_option_.find(op_name);
+      if (it == op_to_option_.end()) {
+        std::cerr << "Didn't find option for  " << op_name << std::endl;
+        fatal = true;
+      } else if (!it->second.empty()) {
+        std::string option_name = it->second;
+
+        if (option_to_struct_.find(option_name) == option_to_struct_.end()) {
+          bool param_struct_found = false;
+          std::string params_guess = std::string("TfLite") + option_name;
+          size_t start = params_guess.find("Options");
+          size_t len = strlen("Options");
+          params_guess.replace(start, len, "Params");
+          for (auto* param = param_structs; *param != nullptr; param++) {
+            if (*param == params_guess) {
+              param_struct_found = true;
+              break;
+            }
+          }
+          if (!param_struct_found) {
+            std::cerr << "Failed to get param struct for option " << option_name
+                      << std::endl;
+          } else {
+            option_to_struct_.insert(std::make_pair(option_name, params_guess));
+          }
+        }
+      }
+    }
+    if (fatal) {
+      exit(1);
+    }
+  }
+
+ private:
+  std::vector<std::string> ops_;
+  std::unordered_map<std::string, std::string> op_to_option_;
+  std::unordered_map<std::string, std::string> option_to_struct_;
+  std::unordered_map<std::string, flatbuffers::TypeFunction>
+      option_to_type_function_;
+};
+
+void GenerateImportForResizeBilinearOp(FILE* fp) {
+  fprintf(fp,
+          "  case BuiltinOperator_RESIZE_BILINEAR:  {\n"
+          "    const auto* params = reinterpret_cast<const "
+          "TfLiteResizeBilinearParams*>(builtin_op_data);\n"
+          "    auto union_type = CreateResizeBilinearOptions(*fbb, "
+          "params->align_corners, params->half_pixel_centers).Union();\n"
+          "    return std::make_pair(BuiltinOptions_ResizeBilinearOptions, "
+          "union_type);\n"
+          "  }\n  break;\n");
+}
+
+void GenerateImportForVarHandleOp(FILE* fp) {
+  fprintf(fp,
+          "  case BuiltinOperator_VAR_HANDLE:  {\n"
+          "    const auto* params = reinterpret_cast<const "
+          "TfLiteVarHandleParams*>(builtin_op_data);\n"
+          "    auto union_type = CreateVarHandleOptions(*fbb, "
+          "fbb->CreateString(params->container), "
+          "fbb->CreateString(params->shared_name)).Union();\n"
+          "    return std::make_pair(BuiltinOptions_VarHandleOptions, "
+          "union_type);\n"
+          "  }\n  break;\n");
+}
+
+// Reshape Op infers output shape either from Parameter or from shape tensor
+// that's is an additional input. When we have this additional shape tensor as
+// input we don't have the parameter present in this layer. In case of more than
+// one input and the shape parameter does not have a valid value, we import an
+// empty vector for the parameters.
+void GenerateImportForReshapeOp(FILE* fp) {
+  fprintf(fp,
+          "  case BuiltinOperator_RESHAPE:  {\n"
+          "    const auto* params = reinterpret_cast<const "
+          "TfLiteReshapeParams*>(builtin_op_data);\n"
+          "    flatbuffers::Offset<void> union_type;\n"
+          "    if (node_inputs_size > 1 && (params->num_dimensions <= 0 || "
+          "params->num_dimensions > TFLITE_RESHAPE_PARAMS_MAX_DIMENSION_COUNT))"
+          " {\n"
+          "      union_type = CreateReshapeOptions(*fbb).Union();\n"
+          "    } else {\n"
+          "      auto val0 = fbb->CreateVector(std::vector<int>(params->shape, "
+          "params->shape + params->num_dimensions));\n"
+          "      union_type = CreateReshapeOptions(*fbb, "
+          "val0).Union();\n"
+          "    }\n"
+          "    return std::make_pair(BuiltinOptions_ReshapeOptions, "
+          "union_type);\n"
+          "  }\n  break;\n");
+}
+
+void GenerateImportForOp(FILE* fp, const std::string& op_name,
+                         const std::string& option_name,
+                         const std::string& option_type,
+                         const flatbuffers::TypeTable* options,
+                         const std::string& struct_name) {
+  // Special-case ResizeBilinear which has some deprecated fields.
+  if (struct_name == "TfLiteResizeBilinearParams") {
+    GenerateImportForResizeBilinearOp(fp);
+    return;
+  }
+
+  if (struct_name == "TfLiteVarHandleParams") {
+    GenerateImportForVarHandleOp(fp);
+    return;
+  }
+
+  // Special case Reshape that may have 'new_shape' field missing from the
+  // parameters.
+  if (struct_name == "TfLiteReshapeParams") {
+    GenerateImportForReshapeOp(fp);
+    return;
+  }
+
+  fprintf(fp, "  case BuiltinOperator_%s:  {\n", op_name.c_str());
+  if (options->num_elems != 0) {
+    fprintf(fp,
+            "    const auto* params = reinterpret_cast<const "
+            "%s*>(builtin_op_data);\n",
+            struct_name.c_str());
+  }
+
+  for (size_t i = 0; i < options->num_elems; i++) {
+    std::string elem_name = options->names[i];
+    bool is_vector = false;
+    std::string vector_name = elem_name;
+    std::string vector_size;
+    std::string vector_type;
+    // TODO(aselle): Irregular naming in builtins
+    if (elem_name == "fused_activation_function")
+      elem_name = "activation";
+    else if (elem_name == "stride_w")
+      elem_name = "stride_width";
+    else if (elem_name == "stride_h")
+      elem_name = "stride_height";
+    else if (elem_name == "stride_d")
+      elem_name = "stride_depth";
+    else if (elem_name == "dilation_h_factor")
+      elem_name = "dilation_height_factor";
+    else if (elem_name == "dilation_w_factor")
+      elem_name = "dilation_width_factor";
+    else if (elem_name == "dilation_d_factor")
+      elem_name = "dilation_depth_factor";
+    else if (elem_name == "idx_out_type")
+      elem_name = "index_out_type";
+
+    // Vector fields treated specially.
+    if (elem_name == "new_shape") {
+      is_vector = true;
+      vector_name = "shape";
+      vector_size = "num_dimensions";
+      vector_type = "int";
+    } else if (elem_name == "squeeze_dims") {
+      is_vector = true;
+      vector_size = "num_squeeze_dims";
+      vector_type = "int";
+    } else if (elem_name == "boundaries") {
+      is_vector = true;
+      vector_size = "num_boundaries";
+      vector_type = "float";
+    } else if (elem_name == "update_window_dims") {
+      is_vector = true;
+      vector_name = "update_window_dims";
+      vector_size = "num_update_window_dims";
+      vector_type = "int64_t";
+    } else if (elem_name == "inserted_window_dims") {
+      is_vector = true;
+      vector_name = "inserted_window_dims";
+      vector_size = "num_inserted_window_dims";
+      vector_type = "int64_t";
+    } else if (elem_name == "scatter_dims_to_operand_dims") {
+      is_vector = true;
+      vector_name = "scatter_dims_to_operand_dims";
+      vector_size = "num_scatter_dims_to_operand_dims";
+      vector_type = "int64_t";
+    }
+
+    if (is_vector) {
+      fprintf(fp,
+              "    auto val%zu = fbb->CreateVector("
+              "std::vector<%s>(params->%s, params->%s + params->%s));\n",
+              i, vector_type.c_str(), vector_name.c_str(), vector_name.c_str(),
+              vector_size.c_str());
+      continue;
+    }
+
+    flatbuffers::TypeCode code = options->type_codes[i];
+    auto contained_type = code.sequence_ref != -1
+                              ? options->type_refs[code.sequence_ref]
+                              : nullptr;
+    std::string mapper = "";
+    if (contained_type == TensorTypeTypeTable) {
+      mapper = "TfLiteTypeToSchemaType";
+    } else if (contained_type == ActivationFunctionTypeTypeTable) {
+      mapper = "TfLiteActivationToSchemaActivation";
+    } else if (contained_type == PaddingTypeTable) {
+      mapper = "TfLitePaddingToSchemaPadding";
+    } else if (contained_type == FullyConnectedOptionsWeightsFormatTypeTable) {
+      mapper = "FullyConnectedOptionsWeightsFormatToSchema";
+    } else if (contained_type == LSTMKernelTypeTypeTable) {
+      mapper = "LSTMKernelTypeToSchema";
+    } else if (contained_type == LSHProjectionTypeTypeTable) {
+      mapper = "LSHProjectionTypeToSchema";
+    } else if (contained_type == MirrorPadModeTypeTable) {
+      mapper = "MirrorPaddingModeToSchema";
+    } else if (contained_type == CombinerTypeTypeTable) {
+      mapper = "CombinerTypeToSchema";
+    }
+
+    fprintf(fp,
+            "    auto val%zu = "
+            "%s(params->%s);\n",
+            i, mapper.c_str(), elem_name.c_str());
+  }
+  fprintf(fp, "    auto union_type = Create%s(*fbb", option_name.c_str());
+  for (size_t i = 0; i < options->num_elems; i++) {
+    fprintf(fp, ", val%zu", i);
+  }
+  fprintf(fp, ").Union();\n");
+  fprintf(fp, "    return std::make_pair(%s, union_type);\n",
+          option_type.c_str());
+  fprintf(fp, "  }\n  break;\n");
+}
+
+void GenerateImport(OpOptionData* option, FILE* fp) {
+  std::unordered_set<std::string> ignores;
+  ignores.insert("CONCAT_EMBEDDINGS");
+  ignores.insert("CALL");
+
+  // Allow any op that doesn't have an options struct to be blocked
+  // together
+  for (const auto& op_name : option->ops()) {
+    auto option_it = option->op_to_option().find(op_name);
+    if (!option_it->second.empty() && ignores.find(op_name) == ignores.end())
+      continue;
+    fprintf(fp, "  case BuiltinOperator_%s:\n", op_name.c_str());
+  }
+  fprintf(fp,
+          "    return std::make_pair(BuiltinOptions_NONE, "
+          "flatbuffers::Offset<void>());\n    break;\n");
+
+  // Iterate over each ops
+  for (const auto& op_name : option->ops()) {
+    if (ignores.find(op_name) != ignores.end()) continue;
+    // Get to the option and struct names, continuing if not found.
+    auto option_it = option->op_to_option().find(op_name);
+    if (option_it->second.empty()) continue;
+    std::string option_name = option_it->second;
+    std::string option_type = "BuiltinOptions_" + option_name;
+    auto option_func_it = option->option_to_type_function().find(option_name);
+    if (option_func_it == option->option_to_type_function().end()) continue;
+    auto struct_name_it = option->option_to_struct().find(option_name);
+    if (struct_name_it == option->option_to_struct().end()) {
+      // If no C struct, then it better have no arguments.
+      auto type_info = option_func_it->second();
+      if (type_info->num_elems != 0) {
+        // We have non-zero arguments in the schema, this means there
+        // should be a struct.
+        fprintf(stderr,
+                "Op %s uses option struct %s which has no builtin struct\n",
+                op_name.c_str(), option_name.c_str());
+        exit(1);
+      }
+      fprintf(fp, "  case BuiltinOperator_%s:\n", op_name.c_str());
+      fprintf(fp, "    return std::make_pair(%s, Create%s(*fbb).Union());",
+              option_type.c_str(), option_name.c_str());
+    } else {
+      // If C struct, then we need to assign all properties
+      auto struct_name = struct_name_it->second;
+      GenerateImportForOp(fp, op_name, option_name, option_type,
+                          option_func_it->second(), struct_name);
+    }
+  }
+  // TODO(aselle): Handle unhandled cases more gracefully.
+  fprintf(fp,
+          "default:    return std::make_pair(BuiltinOptions_NONE, "
+          "flatbuffers::Offset<void>());\n    break;\n");
+}
+
+}  // namespace tflite
+
+int main(int argc, char* argv[]) {
+  tflite::OpOptionData option;
+  if (argc != 2) {
+    fprintf(stderr, "Usage: %s <fname out>\n", argv[0]);
+    return 1;
+  }
+  FILE* fp = fopen(argv[1], "w");
+  tflite::GenerateImport(&option, fp);
+  fclose(fp);
+}

tensorflow/lite/tools/serialization/writer.cc

@@ -0,0 +1,45 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+// Just does a read/write loop of tflite file format using the interpreter as
+// an intermediate.
+//
+// Usage:
+//   writer <input tflite> <output tflite>
+
+#include <cstdio>
+#include <memory>
+
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/tools/serialization/writer_lib.h"
+
+int main(int argc, char* argv[]) {
+  if (argc != 3) {
+    fprintf(stderr, "Usage: %s input_file output_file\n", argv[0]);
+    return 1;
+  }
+  std::unique_ptr<tflite::FlatBufferModel> model =
+      tflite::FlatBufferModel::BuildFromFile(argv[1]);
+  std::unique_ptr<tflite::Interpreter> interpreter;
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates
+      builtin_op_resolver;
+  tflite::InterpreterBuilder(*model, builtin_op_resolver)(&interpreter);
+  tflite::ModelWriter writer(interpreter.get());
+  writer.Write(argv[2]);
+
+  return 0;
+}

tensorflow/lite/tools/serialization/writer_lib.cc

@@ -0,0 +1,609 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+#include "tensorflow/lite/tools/serialization/writer_lib.h"
+
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "flatbuffers/base.h"  // from @flatbuffers
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/string.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/compiler/mlir/lite/schema/schema_conversion_utils.h"
+#include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/interpreter.h"
+#if FLATBUFFERS_LITTLEENDIAN == 0
+#include "tensorflow/lite/core/model_builder.h"
+#endif
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
+#include "tensorflow/compiler/mlir/lite/tools/versioning/op_version.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/tools/serialization/enum_mapping.h"
+#include "tensorflow/lite/version.h"
+
+namespace tflite {
+namespace {
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<OperatorCode>>>
+CreateOpCodeTableImpl(flatbuffers::FlatBufferBuilder* fbb,
+                      std::vector<OpCode>* opcodes) {
+  std::vector<flatbuffers::Offset<OperatorCode>> codes;
+  for (const auto& it : *opcodes) {
+    const char* custom_name = it.custom.empty() ? nullptr : it.custom.c_str();
+    // Use version 0 for builtin op. This is a way to serialize version field to
+    // flatbuffer (since 0 is non default) and it will be corrected later.
+    int32_t op_version = it.builtin != tflite::BuiltinOperator_CUSTOM ? 0 : 1;
+    codes.push_back(
+        CreateOperatorCodeDirect(*fbb, static_cast<BuiltinOperator>(it.builtin),
+                                 custom_name, op_version));
+  }
+  return fbb->template CreateVector<flatbuffers::Offset<OperatorCode>>(codes);
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>>
+ExportBuffersImpl(flatbuffers::FlatBufferBuilder* fbb,
+                  std::vector<std::pair<const uint8_t*, size_t>>* buffers) {
+  std::vector<flatbuffers::Offset<Buffer>> buffer_vector;
+  for (auto buffer : *buffers) {
+    auto data_offset = fbb->CreateVector(buffer.first, buffer.second);
+    buffer_vector.push_back(CreateBuffer(*fbb, data_offset));
+  }
+  return fbb->template CreateVector<flatbuffers::Offset<Buffer>>(buffer_vector);
+}
+
+TfLiteStatus WriteImpl(const std::string& filename, void* data, size_t size) {
+  FILE* fp = fopen(filename.c_str(), "wb");
+  if (!fp) return kTfLiteError;
+
+#if FLATBUFFERS_LITTLEENDIAN == 0
+  const tflite::Model* input_model = tflite::GetModel(data);
+  tflite::FlatBufferModel::ByteSwapTFLiteModel(input_model);
+#endif
+
+  const int result_size = fwrite(data, 1, size, fp);
+  fclose(fp);
+  if (result_size != size) return kTfLiteError;
+
+  return kTfLiteOk;
+}
+
+std::pair<BuiltinOptions, flatbuffers::Offset<void>> CreateBuiltinUnion(
+    flatbuffers::FlatBufferBuilder* fbb, enum BuiltinOperator op,
+    void* builtin_op_data, int node_inputs_size) {
+  switch (op) {
+#include "tensorflow/lite/tools/serialization/option_writer_generated.h"
+  }
+  return std::make_pair(BuiltinOptions_NONE, flatbuffers::Offset<void>());
+}
+
+}  // namespace
+
+template <class T_OUTPUT, class T_INPUT>
+flatbuffers::Offset<flatbuffers::Vector<T_OUTPUT>> SubgraphWriter::ExportVector(
+    flatbuffers::FlatBufferBuilder* fbb, const T_INPUT& v) {
+  std::vector<T_OUTPUT> inputs(v.begin(), v.end());
+  return fbb->template CreateVector<T_OUTPUT>(inputs);
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Operator>>>
+SubgraphWriter::ExportOperators(flatbuffers::FlatBufferBuilder* fbb) {
+  std::vector<flatbuffers::Offset<Operator>> operators;
+
+  std::vector<int> operator_to_opcode;
+  // TODO(aselle): Augment this once we put execution plan in schema.
+  operator_to_opcode.resize(subgraph_->nodes_size(), -1);
+  for (int op_index : execution_plan_) {
+    const auto* node_and_registration =
+        subgraph_->node_and_registration(op_index);
+    const TfLiteRegistration* registration = &node_and_registration->second;
+    if (!registration->custom_name) {
+      operator_to_opcode[op_index] =
+          GetOpCodeForBuiltin(registration->builtin_code);
+    } else {
+      operator_to_opcode[op_index] =
+          GetOpCodeForCustom(registration->custom_name);
+    }
+  }
+  // second pass serialize operators
+  for (int op_index : execution_plan_) {
+    const auto* node_and_registration =
+        subgraph_->node_and_registration(op_index);
+    const TfLiteNode& node = node_and_registration->first;
+    const TfLiteRegistration& registration = node_and_registration->second;
+    flatbuffers::Offset<void> builtin_options;
+    BuiltinOptions builtin_options_type = BuiltinOptions_NONE;
+    // Custom data
+    // TODO(aselle): Custom options format is not known by default. Just assume
+    // for now.
+    auto custom_options_format = CustomOptionsFormat_FLEXBUFFERS;
+    flatbuffers::Offset<flatbuffers::Vector<uint8_t>> custom_options = 0;
+
+    if (!registration.custom_name) {
+      // builtin
+      auto builtin_options_and_type = CreateBuiltinUnion(
+          fbb, static_cast<enum BuiltinOperator>(registration.builtin_code),
+          node.builtin_data, node.inputs->size);
+      builtin_options = builtin_options_and_type.second;
+      builtin_options_type = builtin_options_and_type.first;
+    } else {
+      auto custom_writer = custom_op_to_writer_.find(registration.custom_name);
+      if (custom_writer != custom_op_to_writer_.end() &&
+          custom_writer->second) {
+        // delegate to custom writer if it exists
+        custom_writer->second(fbb, subgraph_, op_index, &custom_options,
+                              &custom_options_format);
+      } else {
+        // use the custom data as fact
+        custom_options = fbb->CreateVector(
+            reinterpret_cast<const uint8_t*>(node.custom_initial_data),
+            node.custom_initial_data_size);
+      }
+    }
+
+    int opcode_index = operator_to_opcode[op_index];
+    std::vector<int> written_inputs =
+        RemapTensorIndicesToWritten(TfLiteIntArrayView(node.inputs));
+    std::vector<int> written_outputs =
+        RemapTensorIndicesToWritten(TfLiteIntArrayView(node.outputs));
+    auto inputs = ExportVector<int32_t>(fbb, written_inputs);
+    auto outputs = ExportVector<int32_t>(fbb, written_outputs);
+    operators.push_back(CreateOperator(*fbb, opcode_index, inputs, outputs,
+                                       builtin_options_type, builtin_options,
+                                       custom_options, custom_options_format));
+  }
+
+  return fbb->template CreateVector<flatbuffers::Offset<Operator>>(operators);
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Tensor>>>
+SubgraphWriter::ExportTensors(flatbuffers::FlatBufferBuilder* fbb) {
+  // Initialized to -1.
+  // A value of -1 means this tensor will not be exported.
+  tensor_to_written_tensor_.resize(subgraph_->tensors_size(), -1);
+
+  std::vector<flatbuffers::Offset<Tensor>> tensors;
+
+  // Make a map from tensor index to whether the tensor is a temporary.
+  std::vector<bool> tensor_is_temporary(subgraph_->tensors_size(), false);
+  for (int op_index = 0; op_index < subgraph_->nodes_size(); ++op_index) {
+    const auto* node_and_registration =
+        subgraph_->node_and_registration(op_index);
+    for (auto tensor_index :
+         TfLiteIntArrayView(node_and_registration->first.temporaries))
+      tensor_is_temporary[tensor_index] = true;
+  }
+
+  // Now we need to remap all used tensor indices
+  int curr_output_index = 0;
+  for (int tensor_index = 0; tensor_index < subgraph_->tensors_size();
+       tensor_index++) {
+    // Temporary tensors and unused tensors will not be written.
+    if (!tensor_is_temporary[tensor_index] &&
+        unused_tensors_.find(tensor_index) == unused_tensors_.end()) {
+      tensor_to_written_tensor_[tensor_index] = curr_output_index++;
+    }
+  }
+
+  for (int tensor_index = 0; tensor_index < subgraph_->tensors_size();
+       ++tensor_index) {
+    // Tensor not exported.
+    if (tensor_to_written_tensor_[tensor_index] == -1) continue;
+
+    if (TfLiteTensor* tensor = subgraph_->tensor(tensor_index)) {
+      // Allocate a buffer index
+      int buffer_index = 0;  // This is null
+      if (tensor->allocation_type == kTfLiteMmapRo) {
+        buffer_index = buffers_->size();
+        buffers_->push_back(std::make_pair(
+            reinterpret_cast<const uint8_t*>(tensor->data.raw), tensor->bytes));
+      }
+      // Primitive type.
+      TensorType type = TfLiteTypeToSchemaType(tensor->type);
+      // Handle quantization
+      flatbuffers::Offset<QuantizationParameters> quantization_params;
+
+      const flatbuffers::Offset<flatbuffers::Vector<float>> null_array;
+      flatbuffers::Offset<flatbuffers::Vector<float>> scale_array;
+      flatbuffers::Offset<flatbuffers::Vector<int64_t>> zero_point_array;
+
+      if (tensor->quantization.type == kTfLiteAffineQuantization) {
+        if (tensor->params.scale != 0.f) {
+          // Quantization with a single argument array.
+          scale_array = fbb->CreateVector<float>({tensor->params.scale});
+          zero_point_array =
+              fbb->CreateVector<int64_t>({tensor->params.zero_point});
+          quantization_params = CreateQuantizationParameters(
+              *fbb, null_array, null_array, scale_array, zero_point_array);
+        } else {  // Multi channel quantization.
+          const TfLiteAffineQuantization* params =
+              reinterpret_cast<TfLiteAffineQuantization*>(
+                  tensor->quantization.params);
+          const size_t num_scales = params->scale->size;
+
+          std::vector<float> scale_vector(params->scale->data,
+                                          params->scale->data + num_scales);
+          // Copy zero point by default.
+          std::vector<int64_t> zero_point_vector(
+              params->zero_point->data,
+              params->zero_point->data + params->zero_point->size);
+          // If we have more zero points, copy them.
+          if (params->zero_point->size != params->scale->size) {
+            zero_point_vector.resize(params->scale->size, zero_point_vector[0]);
+          }
+          scale_array = fbb->CreateVector<float>(scale_vector);
+          zero_point_array = fbb->CreateVector<int64_t>(zero_point_vector);
+          quantization_params = CreateQuantizationParameters(
+              *fbb, null_array, null_array, scale_array, zero_point_array,
+              QuantizationDetails_NONE, 0, params->quantized_dimension);
+        }
+      }
+
+      // Shape
+      // Some tensors added during op init are not registered formally as
+      // node temporaries. Some didn't get memory allocated for them, and we
+      // should avoid serializing those tensors.
+      if (tensor->dims) {
+        TfLiteIntArrayView shape_view(tensor->dims);
+        std::vector<int> shape =
+            std::vector<int>(shape_view.begin(), shape_view.end());
+
+        Offset<flatbuffers::String> tensor_name_offset = 0;
+        if (tensor->name != nullptr) {
+          tensor_name_offset = fbb->CreateString(tensor->name);
+        }
+
+        flatbuffers::Offset<flatbuffers::Vector<int32_t>>
+            shape_signature_offset = 0;
+
+        if (serialize_dims_signature_ && tensor->dims_signature != nullptr) {
+          TfLiteIntArrayView shape_signature_view(tensor->dims_signature);
+          std::vector<int32_t> shape_signature(shape_signature_view.begin(),
+                                               shape_signature_view.end());
+          shape_signature_offset = ExportVector<int32_t>(fbb, shape_signature);
+        }
+
+        // TFLite runtime does not differentiate between unranked and scalar
+        // tensors. Assume shapeless tensors are scalars when serializing.
+        // TODO(b/255826755): Remove workaround when runtime can differentiate
+        // between scalar and unranked tensors.
+        bool has_rank = true;
+        tensors.push_back(CreateTensor(
+            *fbb, ExportVector<int32_t>(fbb, shape), type, buffer_index,
+            tensor_name_offset, quantization_params, tensor->is_variable,
+            /*sparsity=*/0, shape_signature_offset, has_rank));
+      }
+    }
+  }
+  return fbb->template CreateVector<flatbuffers::Offset<Tensor>>(tensors);
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>>
+SubgraphWriter::ExportBuffers(flatbuffers::FlatBufferBuilder* fbb) {
+  return ExportBuffersImpl(fbb, buffers_);
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<OperatorCode>>>
+SubgraphWriter::CreateOpCodeTable(flatbuffers::FlatBufferBuilder* fbb) {
+  return CreateOpCodeTableImpl(fbb, opcodes_);
+}
+
+template <class T>
+std::vector<int> SubgraphWriter::RemapTensorIndicesToWritten(const T& input) {
+  std::vector<int> output;
+  output.reserve(input.size());
+  for (int x : input) {
+    // Special value representing an optional tensor which is not present.
+    if (x == -1) {
+      output.push_back(x);
+      continue;
+    }
+    if (tensor_to_written_tensor_[x] != -1) {
+      output.push_back(tensor_to_written_tensor_[x]);
+    }
+  }
+  return output;
+}
+
+TfLiteStatus SubgraphWriter::GetBuffer(std::unique_ptr<uint8_t[]>* out,
+                                       size_t* size) {
+  if (!out || !size) return kTfLiteError;
+  flatbuffers::FlatBufferBuilder builder(/*initial_size=*/10240);
+  std::vector<flatbuffers::Offset<SubGraph>> subgraphs_as_vector;
+  subgraphs_as_vector.push_back(
+      PopulateAndGetOffset(&builder, subgraph_->GetName()));
+
+  flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>>
+      buffers = ExportBuffers(&builder);
+
+  auto description = builder.CreateString("Exported from Subgraph.");
+
+  auto op_codes = CreateOpCodeTable(&builder);
+  auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, op_codes,
+                           builder.CreateVector(subgraphs_as_vector),
+                           description, buffers);
+  ::tflite::FinishModelBuffer(builder, model);
+  ::tflite::UpdateOpVersion(builder.GetBufferPointer());
+  const uint8_t* buffer = builder.GetBufferPointer();
+  *size = builder.GetSize();
+  (*out).reset(new uint8_t[*size]);
+  memcpy(out->get(), buffer, *size);
+  return kTfLiteOk;
+}
+
+flatbuffers::Offset<SubGraph> SubgraphWriter::PopulateAndGetOffset(
+    flatbuffers::FlatBufferBuilder* builder, const std::string& subgraph_name) {
+  auto tensors = ExportTensors(builder);
+  std::vector<int> written_inputs = RemapTensorIndicesToWritten(inputs_);
+  std::vector<int> written_outputs = RemapTensorIndicesToWritten(outputs_);
+  auto inputs = ExportVector<int32_t>(builder, written_inputs);
+  auto outputs = ExportVector<int32_t>(builder, written_outputs);
+
+  auto ops = ExportOperators(builder);
+  auto name = builder->CreateString(subgraph_name);
+  return CreateSubGraph(*builder, tensors, inputs, outputs, ops, name);
+}
+
+TfLiteStatus SubgraphWriter::Write(const std::string& filename) {
+  std::unique_ptr<uint8_t[]> buffer;
+  size_t size;
+  TF_LITE_ENSURE_STATUS(GetBuffer(&buffer, &size));
+  return WriteImpl(filename, buffer.get(), size);
+}
+
+TfLiteStatus SubgraphWriter::RegisterCustomWriter(
+    const std::string& custom_name, CustomWriter custom_writer) {
+  if (custom_op_to_writer_.find(custom_name) != custom_op_to_writer_.end()) {
+    return kTfLiteError;
+  }
+  custom_op_to_writer_.insert(std::make_pair(custom_name, custom_writer));
+  return kTfLiteOk;
+}
+
+TfLiteStatus SubgraphWriter::CheckInputOutput(
+    const std::vector<int>& inputs, const std::vector<int>& outputs,
+    const std::vector<int>& execution_plan) {
+  absl::flat_hash_set<int> known_tensors(inputs.begin(), inputs.end());
+  known_tensors.insert(subgraph_->variables().begin(),
+                       subgraph_->variables().end());
+  // Scan execution plan and confirm input tensors are known before each node
+  // executes. Then append output tensors to known tensors.
+  for (int op_index : execution_plan) {
+    const auto* node_and_registration =
+        subgraph_->node_and_registration(op_index);
+    const TfLiteNode& node = node_and_registration->first;
+    for (int tensor_index : TfLiteIntArrayView(node.inputs)) {
+      if (tensor_index < 0) {
+        // Skip if optional input not present.
+        if (tensor_index == kTfLiteOptionalTensor) {
+          continue;
+        } else {
+          return kTfLiteError;
+        }
+      }
+      if (TfLiteTensor* tensor = subgraph_->tensor(tensor_index)) {
+        // Skip constant tensors.
+        if (tensor->allocation_type == kTfLiteMmapRo) {
+          continue;
+        }
+      }
+
+      if (known_tensors.find(tensor_index) == known_tensors.end()) {
+        subgraph_->context()->ReportError(
+            subgraph_->context(),
+            "Node (%d) uses an input (%d) that is not provided.", op_index,
+            tensor_index);
+        return kTfLiteError;
+      }
+    }
+    TfLiteIntArrayView outputs(node.outputs);
+    known_tensors.insert(outputs.begin(), outputs.end());
+  }
+
+  // Check if outputs are known tensors or constants.
+  for (int tensor_index : outputs) {
+    if (TfLiteTensor* tensor = subgraph_->tensor(tensor_index)) {
+      // Skip constant tensors.
+      if (tensor->allocation_type == kTfLiteMmapRo) {
+        continue;
+      }
+    }
+
+    if (known_tensors.find(tensor_index) == known_tensors.end()) {
+      subgraph_->context()->ReportError(
+          subgraph_->context(),
+          "Output (%d) is not produced by the execution plan.", tensor_index);
+      return kTfLiteError;
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus SubgraphWriter::SetCustomInputOutput(
+    const std::vector<int>& inputs, const std::vector<int>& outputs,
+    const std::vector<int>& execution_plan) {
+  TF_LITE_ENSURE_STATUS(CheckInputOutput(inputs, outputs, execution_plan));
+  inputs_ = inputs;
+  outputs_ = outputs;
+  execution_plan_ = execution_plan;
+  return kTfLiteOk;
+}
+
+ModelWriter::ModelWriter(Interpreter* interpreter,
+                         bool serialize_dims_signature) {
+  std::vector<Subgraph*> subgraphs;
+
+  // Retrieves the list of the subgraphs from the interpreter for constructing
+  // a list of SubgraphWriters.
+  subgraphs.reserve(interpreter->subgraphs_size());
+  for (int i = 0; i < interpreter->subgraphs_size(); ++i) {
+    subgraphs.push_back(interpreter->subgraph(i));
+  }
+
+  Init(subgraphs, serialize_dims_signature);
+}
+
+ModelWriter::ModelWriter(const std::vector<Subgraph*>& subgraphs,
+                         bool serialize_dims_signature) {
+  Init(subgraphs, serialize_dims_signature);
+}
+
+void ModelWriter::Init(const std::vector<Subgraph*>& subgraphs,
+                       bool serialize_dims_signature) {
+  buffers_.push_back(std::make_pair(nullptr, 0));
+  subgraph_writers_.reserve(subgraphs.size());
+  for (auto* subgraph : subgraphs) {
+    SubgraphWriter writer(subgraph, &buffers_, &opcodes_,
+                          &builtin_op_to_opcode_, serialize_dims_signature);
+    subgraph_writers_.push_back(writer);
+  }
+
+  // Populate subgraph_index_mapper_.
+  if (!subgraphs.empty()) {
+    absl::flat_hash_map<Subgraph*, int> subgraph_to_new_subgraph_index;
+    for (int i = 0; i < subgraphs.size(); ++i) {
+      subgraph_to_new_subgraph_index[subgraphs[i]] = i;
+    }
+
+    auto* all_subgraphs = subgraphs[0]->GetSubgraphs();
+    for (int i = 0; i < all_subgraphs->size(); ++i) {
+      auto it = subgraph_to_new_subgraph_index.find(all_subgraphs->at(i));
+      if (it != subgraph_to_new_subgraph_index.end()) {
+        subgraph_index_mapper_[i] = it->second;
+      }
+    }
+  }
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>>
+ModelWriter::ExportBuffers(flatbuffers::FlatBufferBuilder* fbb) {
+  return ExportBuffersImpl(fbb, &buffers_);
+}
+
+flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<OperatorCode>>>
+ModelWriter::CreateOpCodeTable(flatbuffers::FlatBufferBuilder* fbb) {
+  return CreateOpCodeTableImpl(fbb, &opcodes_);
+}
+
+TfLiteStatus ModelWriter::GetBuffer(std::unique_ptr<uint8_t[]>* out,
+                                    size_t* size) {
+  if (!out || !size) return kTfLiteError;
+  flatbuffers::FlatBufferBuilder builder(/*initial_size=*/10240);
+
+  std::vector<flatbuffers::Offset<SubGraph>> subgraphs_as_vector;
+  subgraphs_as_vector.reserve(subgraph_writers_.size());
+  for (auto& subgraph_writer : subgraph_writers_) {
+    subgraphs_as_vector.push_back(subgraph_writer.PopulateAndGetOffset(
+        &builder, subgraph_writer.subgraph_->GetName()));
+  }
+
+  flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<Buffer>>>
+      buffers = ExportBuffers(&builder);
+
+  auto description = builder.CreateString("Exported from Subgraph.");
+
+  auto op_codes = CreateOpCodeTable(&builder);
+  auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, op_codes,
+                           builder.CreateVector(subgraphs_as_vector),
+                           description, buffers);
+  ::tflite::FinishModelBuffer(builder, model);
+  ::tflite::UpdateOpVersion(builder.GetBufferPointer());
+  UpdateSubgraphReferences(&builder);
+  const uint8_t* buffer = builder.GetBufferPointer();
+  *size = builder.GetSize();
+  (*out).reset(new uint8_t[*size]);
+  memcpy(out->get(), buffer, *size);
+  return kTfLiteOk;
+}
+
+TfLiteStatus ModelWriter::Write(const std::string& filename) {
+  std::unique_ptr<uint8_t[]> buffer;
+  size_t size;
+  TF_LITE_ENSURE_STATUS(GetBuffer(&buffer, &size));
+  return WriteImpl(filename, buffer.get(), size);
+}
+
+void ModelWriter::SetUnusedTensors(int subgraph_index,
+                                   const std::set<int>& unused_tensors) {
+  subgraph_writers_[subgraph_index].SetUnusedTensors(unused_tensors);
+}
+
+TfLiteStatus ModelWriter::SetCustomInputOutput(
+    int subgraph_index, const std::vector<int>& inputs,
+    const std::vector<int>& outputs, const std::vector<int>& execution_plan) {
+  return subgraph_writers_[subgraph_index].SetCustomInputOutput(inputs, outputs,
+                                                                execution_plan);
+}
+
+TfLiteStatus ModelWriter::RegisterCustomWriter(const std::string& custom_name,
+                                               CustomWriter custom_writer) {
+  for (auto& subgraph_writer : subgraph_writers_) {
+    subgraph_writer.RegisterCustomWriter(custom_name, custom_writer);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus ModelWriter::UpdateSubgraphReferences(
+    flatbuffers::FlatBufferBuilder* fbb) {
+  auto model = tflite::GetMutableModel(fbb->GetBufferPointer());
+
+  for (SubGraph* subgraph : *model->mutable_subgraphs()) {
+    for (Operator* op : *subgraph->mutable_operators()) {
+      if (op->builtin_options_type() == BuiltinOptions_WhileOptions) {
+        auto while_options =
+            static_cast<tflite::WhileOptions*>(op->mutable_builtin_options());
+        auto new_cond_index =
+            subgraph_index_mapper_.find(while_options->cond_subgraph_index());
+        auto new_body_index =
+            subgraph_index_mapper_.find(while_options->body_subgraph_index());
+        if (new_cond_index == subgraph_index_mapper_.end() ||
+            new_body_index == subgraph_index_mapper_.end()) {
+          // Subgraph not found in the map.
+          return kTfLiteError;
+        }
+        while_options->mutate_cond_subgraph_index(new_cond_index->second);
+        while_options->mutate_body_subgraph_index(new_body_index->second);
+      } else if (op->builtin_options_type() == BuiltinOptions_IfOptions) {
+        auto if_options =
+            static_cast<tflite::IfOptions*>(op->mutable_builtin_options());
+        auto new_then_index =
+            subgraph_index_mapper_.find(if_options->then_subgraph_index());
+        auto new_else_index =
+            subgraph_index_mapper_.find(if_options->else_subgraph_index());
+        if (new_then_index == subgraph_index_mapper_.end() ||
+            new_else_index == subgraph_index_mapper_.end()) {
+          // Subgraph not found in the map.
+          return kTfLiteError;
+        }
+        if_options->mutate_then_subgraph_index(new_then_index->second);
+        if_options->mutate_else_subgraph_index(new_else_index->second);
+      }
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace tflite

tensorflow/lite/tools/serialization/writer_lib.h

@@ -0,0 +1,292 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+// Library to write a flatbuffer of a currently loaded TFLite model/subgraph.
+
+#ifndef TENSORFLOW_LITE_TOOLS_SERIALIZATION_WRITER_LIB_H_
+#define TENSORFLOW_LITE_TOOLS_SERIALIZATION_WRITER_LIB_H_
+#include <cstddef>
+#include <cstdint>
+#include <iostream>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+
+// This #include needs to precede the inclusion of any other TF Lite header
+// file that might depend on the non-mutable schema_generated.h, directly,
+// e.g. core/api/op_resolver.h, or indirectly, e.g. core/subgraph.h.
+// That's because "tensorflow/lite/schema/mutable/schema_generated.h"
+// and "tensorflow/lite/schema/schema_generated.h" both use the same
+// header guard macro (FLATBUFFERS_GENERATED_SCHEMA_TFLITE_H_), but have
+// different contents (the former is a superset of the latter). In particular
+// the one in mutable/ is built with the "--gen-mutable" and "--gen-object-api"
+// flags to the flatbuffer schema compiler which cause some additional
+// (non-virtual) accessor methods and API functions to be declared.
+// The code here uses those methods, so we need to make sure that we get
+// the mutable variant of this header.
+//
+//  The '#if' here prevents automatic reordering of this #include.
+#if 1
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
+#endif
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/lite/builtin_op_data.h"
+#include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/tools/serialization/enum_mapping.h"
+#include "tensorflow/lite/version.h"
+
+namespace tflite {
+
+struct OpCode {
+  int builtin;
+  std::string custom;
+};
+
+// Forward declaration.
+class SubgraphWriter;
+
+// Handles writing a full TFLite model (with 1 or more subgraphs) to a
+// serialized TF lite file format.
+// TODO(b/174708523): Support custom I/O or unused tensors later.
+class ModelWriter {
+ public:
+  // CustomWriter allows the delegate to customize the write to the flatbuffer.
+  typedef flatbuffers::Offset<Operator> (*CustomWriter)(
+      flatbuffers::FlatBufferBuilder* fbb, Subgraph* subgraph, int node_index,
+      flatbuffers::Offset<flatbuffers::Vector<uint8_t>>* output_options,
+      CustomOptionsFormat* custom_options_format);
+
+  // Construct a writer for the specified `interpreter`. Then, use
+  // .Write() or .GetBuffer(...) to extract the data.
+  explicit ModelWriter(Interpreter* interpreter,
+                       bool serialize_dims_signature = true);
+
+  // Same as above, except takes subgraphs as input.
+  explicit ModelWriter(const std::vector<Subgraph*>& subgraphs,
+                       bool serialize_dims_signature = true);
+
+  // Get a buffer and size of a serialized flatbuffer.
+  TfLiteStatus GetBuffer(std::unique_ptr<uint8_t[]>* out, size_t* size);
+  // Write the serialized flatbuffer to the prescribed `filename`.
+  TfLiteStatus Write(const std::string& filename);
+
+  // Specifies unused tensors on the target subgraph.
+  void SetUnusedTensors(int subgraph_index,
+                        const std::set<int>& unused_tensors);
+
+  // Specifies custom inputs, outputs, and execution_plan to target subgraph.
+  TfLiteStatus SetCustomInputOutput(int subgraph_index,
+                                    const std::vector<int>& inputs,
+                                    const std::vector<int>& outputs,
+                                    const std::vector<int>& execution_plan);
+
+  // Registers a custom writer for a custom op. The customization allows the
+  // caller to change the custom data.
+  TfLiteStatus RegisterCustomWriter(const std::string& custom_name,
+                                    CustomWriter custom_writer);
+
+ private:
+  // For initializing the ModelWriter internal data.
+  void Init(const std::vector<Subgraph*>& subgraphs,
+            bool serialize_dims_signature);
+
+  template <class T>
+  using Offset = flatbuffers::Offset<T>;
+  Offset<flatbuffers::Vector<Offset<OperatorCode>>> CreateOpCodeTable(
+      flatbuffers::FlatBufferBuilder* fbb);
+  Offset<flatbuffers::Vector<Offset<Buffer>>> ExportBuffers(
+      flatbuffers::FlatBufferBuilder* fbb);
+
+  // Updates subgraph_index references of control flow ops in the serialized
+  // model. This is necessary since the order of the subgraphs in the serialized
+  // model might be different than the original input model.
+  TfLiteStatus UpdateSubgraphReferences(flatbuffers::FlatBufferBuilder* fbb);
+
+  // List of subgraph writers owned by this model writer.
+  // There is one subgraph writer for each subgraph in the model.
+  std::vector<SubgraphWriter> subgraph_writers_;
+
+  // This data corresponds to the overall model (rather than individual
+  // subgraphs), so we define common fields. Keep track of byte buffers
+  std::vector<std::pair<const uint8_t*, size_t>> buffers_;
+  // List of used opcodes
+  std::vector<OpCode> opcodes_;
+  absl::flat_hash_map<int, int> builtin_op_to_opcode_;
+
+  // Map from original subgraph indices to the new indices.
+  absl::flat_hash_map<int, int> subgraph_index_mapper_;
+};
+
+// Handles writing TensorFlow Lite running subgraph to a serialized TF lite
+// file format.
+// TODO(b/174708523): Reconcile into ModelWriter?
+class SubgraphWriter {
+ public:
+  friend class ModelWriter;
+
+  typedef flatbuffers::Offset<Operator> (*CustomWriter)(
+      flatbuffers::FlatBufferBuilder* fbb, Subgraph* subgraph, int node_index,
+      flatbuffers::Offset<flatbuffers::Vector<uint8_t>>* output_options,
+      CustomOptionsFormat* custom_options_format);
+
+  // Construct a subgraph writer for the specified `subgraph`. Then, use
+  // .Write() or .GetBuffer(...) to extract the data.
+  explicit SubgraphWriter(Subgraph* subgraph,
+                          bool serialize_dims_signature = true)
+      : subgraph_(subgraph),
+        inputs_(subgraph->inputs()),
+        outputs_(subgraph->outputs()),
+        execution_plan_(subgraph->execution_plan()),
+        serialize_dims_signature_(serialize_dims_signature) {
+    buffers_ = &buffers_data_;
+    opcodes_ = &opcodes_data_;
+    builtin_op_to_opcode_ = &builtin_op_to_opcode_data_;
+    buffers_->push_back(std::make_pair(nullptr, 0));
+  }
+
+  // Get a buffer and size of a serialized flatbuffer.
+  TfLiteStatus GetBuffer(std::unique_ptr<uint8_t[]>* out, size_t* size);
+  // Write the serialized flatbuffer to the prescribed `filename`.
+  TfLiteStatus Write(const std::string& filename);
+  // Registers a custom writer for a custom op. The customization allows the
+  // caller to change the custom data.
+  TfLiteStatus RegisterCustomWriter(const std::string& custom_name,
+                                    CustomWriter custom_writer);
+  // Tensors that are unused and shouldn't be written.
+  void SetUnusedTensors(const std::set<int>& unused_tensors) {
+    unused_tensors_ = unused_tensors;
+  }
+  // Sets custom inputs, outputs, and execution_plan so that a portion of the
+  // subgraph is written to the buffer instead of the whole subgraph.
+  TfLiteStatus SetCustomInputOutput(const std::vector<int>& inputs,
+                                    const std::vector<int>& outputs,
+                                    const std::vector<int>& execution_plan);
+
+ private:
+  // Used by ModelWriter.
+  explicit SubgraphWriter(
+      Subgraph* subgraph,
+      std::vector<std::pair<const uint8_t*, size_t>>* external_buffers,
+      std::vector<OpCode>* external_opcodes,
+      absl::flat_hash_map<int, int>* external_builtin_op_to_opcode,
+      bool serialize_dims_signature)
+      : subgraph_(subgraph),
+        inputs_(subgraph->inputs()),
+        outputs_(subgraph->outputs()),
+        execution_plan_(subgraph->execution_plan()),
+        serialize_dims_signature_(serialize_dims_signature) {
+    buffers_ = external_buffers;
+    opcodes_ = external_opcodes;
+    builtin_op_to_opcode_ = external_builtin_op_to_opcode;
+    buffers_->push_back(std::make_pair(nullptr, 0));
+  }
+
+  // Used by ModelWriter to populate data specific to this subgraph.
+  // Global stuff (like opcodes & buffers) is populated into buffers_, opcodes_,
+  // etc. & populated in the Flatbuffer by ModelWriter.
+  flatbuffers::Offset<SubGraph> PopulateAndGetOffset(
+      flatbuffers::FlatBufferBuilder* builder,
+      const std::string& subgraph_name);
+
+  template <class T>
+  using Offset = flatbuffers::Offset<T>;
+  template <class T_OUTPUT, class T_INPUT>
+  Offset<flatbuffers::Vector<T_OUTPUT>> ExportVector(
+      flatbuffers::FlatBufferBuilder* fbb, const T_INPUT& v);
+  Offset<flatbuffers::Vector<Offset<Tensor>>> ExportTensors(
+      flatbuffers::FlatBufferBuilder* fbb);
+  Offset<flatbuffers::Vector<Offset<Operator>>> ExportOperators(
+      flatbuffers::FlatBufferBuilder* fbb);
+  Offset<flatbuffers::Vector<Offset<OperatorCode>>> CreateOpCodeTable(
+      flatbuffers::FlatBufferBuilder* fbb);
+  Offset<flatbuffers::Vector<Offset<Buffer>>> ExportBuffers(
+      flatbuffers::FlatBufferBuilder* fbb);
+
+  template <class T>
+  std::vector<int> RemapTensorIndicesToWritten(const T& input);
+
+  // Checks if given `input`, `output`, and `execution_plan` represents a valid
+  // model within the Subgraph.
+  TfLiteStatus CheckInputOutput(const std::vector<int>& inputs,
+                                const std::vector<int>& outputs,
+                                const std::vector<int>& execution_plan);
+
+  int GetOpCodeForBuiltin(int builtin_op_index) {
+    // auto it = builtin_op_to_opcode_.find(builtin_op_index);
+    std::pair<decltype(builtin_op_to_opcode_data_)::iterator, bool> result =
+        builtin_op_to_opcode_->insert(
+            std::make_pair(builtin_op_index, opcodes_->size()));
+    if (result.second) {
+      opcodes_->push_back({builtin_op_index, ""});
+    }
+    return result.first->second;
+  }
+
+  int GetOpCodeForCustom(const std::string& custom_name) {
+    std::pair<decltype(custom_op_to_opcode_)::iterator, bool> result =
+        custom_op_to_opcode_.insert(
+            std::make_pair(custom_name, opcodes_->size()));
+    if (result.second) {
+      opcodes_->push_back({BuiltinOperator_CUSTOM, custom_name});
+    }
+    return result.first->second;
+  }
+
+  // The subgraph we are writing
+  Subgraph* subgraph_;
+  // Input tensor indices to be written.
+  std::vector<int> inputs_;
+  // Output tensor indices to be written.
+  std::vector<int> outputs_;
+  // Order of nodes to be written.
+  std::vector<int> execution_plan_;
+  // List of op codes and mappings from builtin or custom op to opcode
+  std::set<int> unused_tensors_;
+  // For every tensor index in the subgraph, the index in the written.
+  // This is different due to temporary and unused tensors not being written.
+  std::vector<int> tensor_to_written_tensor_;
+  std::unordered_map<std::string, int> custom_op_to_opcode_;
+  std::unordered_map<std::string, CustomWriter> custom_op_to_writer_;
+
+  // We use pointers for these, since they may be provided by ModelWriter.
+  // Keep track of byte buffers
+  std::vector<std::pair<const uint8_t*, size_t>>* buffers_;
+  // List of used opcodes
+  std::vector<OpCode>* opcodes_;
+  absl::flat_hash_map<int, int>* builtin_op_to_opcode_;
+
+  // These are used if SubgraphWriter is being used directly.
+  std::vector<std::pair<const uint8_t*, size_t>> buffers_data_;
+  // List of used opcodes
+  std::vector<OpCode> opcodes_data_;
+  absl::flat_hash_map<int, int> builtin_op_to_opcode_data_;
+
+  // Specifies whether tensor dims_signature should be serialized.
+  bool serialize_dims_signature_;
+};
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_SERIALIZATION_WRITER_LIB_H_

tensorflow/lite/tools/serialization/writer_lib_test.cc

@@ -0,0 +1,755 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+#include "tensorflow/lite/tools/serialization/writer_lib.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <ios>
+#include <memory>
+#include <numeric>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/log/check.h"
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
+#include "tensorflow/lite/builtin_ops.h"
+#include "tensorflow/lite/context_util.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/kernels/builtin_op_kernels.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/kernels/subgraph_test_util.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/util.h"
+#include "tsl/platform/logging.h"
+
+namespace tflite {
+
+using subgraph_test_util::CheckIntTensor;
+using subgraph_test_util::FillIntTensor;
+
+std::string CreateFilePath(const std::string& file_name) {
+  const char* tmp_dir = getenv("TEST_TMPDIR");
+  return std::string(tmp_dir ? tmp_dir : "./") + file_name;
+}
+
+// The bool param indicates whether we use SubgraphWriter(true) or
+// ModelWriter(false) for the test
+class SingleSubgraphTest : public ::testing::TestWithParam<bool> {
+ protected:
+  void WriteToFile(Interpreter* interpreter, const std::string& filename,
+                   bool use_subgraph_writer) {
+    if (use_subgraph_writer) {
+      SubgraphWriter writer(&interpreter->primary_subgraph());
+      CHECK_EQ(writer.Write(filename), kTfLiteOk);
+    } else {
+      ModelWriter writer(interpreter);
+      CHECK_EQ(writer.Write(filename), kTfLiteOk);
+    }
+  }
+};
+
+TEST_P(SingleSubgraphTest, InvalidDestinations) {
+  Interpreter interpreter;
+  interpreter.AddTensors(3);
+  float foo[] = {1, 2, 3};
+  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadOnly(
+      1, kTfLiteFloat32, "b", {3}, TfLiteQuantization(),
+      reinterpret_cast<char*>(foo), sizeof(foo));
+  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {3},
+                                           TfLiteQuantization());
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({2});
+  const char* initial_data = "";
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+  interpreter.AddNodeWithParameters({0, 1}, {2}, initial_data, 0,
+                                    reinterpret_cast<void*>(builtin_data), reg);
+
+  // Check if invalid filename is handled gracefully.
+  if (GetParam()) {
+    SubgraphWriter writer(&interpreter.primary_subgraph());
+    CHECK_EQ(writer.Write(""), kTfLiteError);
+  } else {
+    ModelWriter writer(&interpreter);
+    CHECK_EQ(writer.Write(""), kTfLiteError);
+  }
+
+  // Check if invalid buffer is handled gracefully.
+  size_t size;
+  if (GetParam()) {
+    SubgraphWriter writer(&interpreter.primary_subgraph());
+    CHECK_EQ(writer.GetBuffer(nullptr, &size), kTfLiteError);
+  } else {
+    ModelWriter writer(&interpreter);
+    CHECK_EQ(writer.GetBuffer(nullptr, &size), kTfLiteError);
+  }
+}
+
+TEST_P(SingleSubgraphTest, FloatModelTest) {
+  Interpreter interpreter;
+  interpreter.AddTensors(3);
+  float foo[] = {1, 2, 3};
+  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadOnly(
+      1, kTfLiteFloat32, "b", {3}, TfLiteQuantization(),
+      reinterpret_cast<char*>(foo), sizeof(foo));
+  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {3},
+                                           TfLiteQuantization());
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({2});
+  const char* initial_data = "";
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+  interpreter.AddNodeWithParameters({0, 1}, {2}, initial_data, 0,
+                                    reinterpret_cast<void*>(builtin_data), reg);
+
+  const std::string test_file = CreateFilePath("test_float.tflite");
+  WriteToFile(&interpreter, test_file, GetParam());
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  CHECK_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+}
+
+// Tests writing only a portion of the subgraph.
+TEST_P(SingleSubgraphTest, CustomInputOutputTest) {
+  Interpreter interpreter;
+  interpreter.AddTensors(4);
+  constexpr float kFoo[] = {1, 2, 3};
+  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadOnly(
+      1, kTfLiteFloat32, "b", {3}, TfLiteQuantization(),
+      reinterpret_cast<const char*>(kFoo), sizeof(kFoo));
+  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadWrite(3, kTfLiteFloat32, "d", {3},
+                                           TfLiteQuantization());
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({3});
+
+  // Add two ops: Add and Relu
+  const char* initial_data = "";
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+  interpreter.AddNodeWithParameters({0, 1}, {2}, initial_data, 0,
+                                    reinterpret_cast<void*>(builtin_data), reg);
+
+  const TfLiteRegistration* reg2 = resolver.FindOp(BuiltinOperator_RELU, 1);
+  interpreter.AddNodeWithParameters({2}, {3}, nullptr, 0, nullptr, reg2);
+
+  // Only write the second op.
+  const std::string test_file = CreateFilePath("test_custom.tflite");
+  SubgraphWriter writer(&interpreter.primary_subgraph());
+  EXPECT_EQ(writer.SetCustomInputOutput(/*inputs=*/{2}, /*outputs=*/{3},
+                                        /*execution_plan=*/{1}),
+            kTfLiteOk);
+  writer.SetUnusedTensors({0, 1});
+  writer.Write(test_file);
+
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  ASSERT_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+}
+
+TEST_P(SingleSubgraphTest, CustomInputOutputErrorCasesTest) {
+  Interpreter interpreter;
+  interpreter.AddTensors(5);
+  constexpr float kFoo[] = {1, 2, 3};
+  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadOnly(
+      1, kTfLiteFloat32, "b", {3}, TfLiteQuantization(),
+      reinterpret_cast<const char*>(kFoo), sizeof(kFoo));
+  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadWrite(3, kTfLiteFloat32, "d", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadWrite(4, kTfLiteFloat32, "e", {3},
+                                           TfLiteQuantization());
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({4});
+
+  // Add three ops.
+  const char* initial_data = "";
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+  interpreter.AddNodeWithParameters({0, 1}, {2}, initial_data, 0,
+                                    reinterpret_cast<void*>(builtin_data), reg);
+
+  const TfLiteRegistration* reg2 = resolver.FindOp(BuiltinOperator_RELU, 1);
+  interpreter.AddNodeWithParameters({2}, {3}, nullptr, 0, nullptr, reg2);
+
+  const TfLiteRegistration* reg3 = resolver.FindOp(BuiltinOperator_RELU6, 1);
+  interpreter.AddNodeWithParameters({3}, {4}, nullptr, 0, nullptr, reg3);
+
+  SubgraphWriter writer(&interpreter.primary_subgraph());
+
+  // Test wrong input.
+  EXPECT_EQ(writer.SetCustomInputOutput(/*inputs=*/{2}, /*outputs=*/{3},
+                                        /*execution_plan=*/{0, 1}),
+            kTfLiteError);
+  // Test wrong output.
+  EXPECT_EQ(writer.SetCustomInputOutput(/*inputs=*/{0, 1}, /*outputs=*/{4},
+                                        /*execution_plan=*/{0, 1}),
+            kTfLiteError);
+  // Test a valid case.
+  EXPECT_EQ(writer.SetCustomInputOutput(/*inputs=*/{0, 1}, /*outputs=*/{3},
+                                        /*execution_plan=*/{0, 1}),
+            kTfLiteOk);
+}
+
+// Tests if SetCustomInputOutput handles variable tensors correctly.
+TEST_P(SingleSubgraphTest, CustomInputOutputVariableTensorTest) {
+  Interpreter interpreter;
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+
+  // Create tensors.
+  interpreter.AddTensors(3);
+  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {3},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadWrite(1, kTfLiteFloat32, "b", {3},
+                                           TfLiteQuantization(),
+                                           /*is_variable=*/true);
+  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {3},
+                                           TfLiteQuantization());
+  interpreter.SetInputs({0});
+  interpreter.SetOutputs({2});
+  interpreter.SetVariables({1});
+
+  // Create an Add node.
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  interpreter.AddNodeWithParameters({0, 1}, {2}, nullptr, 0,
+                                    reinterpret_cast<void*>(builtin_data),
+                                    resolver.FindOp(BuiltinOperator_ADD, 1));
+
+  // Write model to file.
+  const std::string test_file = CreateFilePath("test_variables.tflite");
+  SubgraphWriter writer(&interpreter.primary_subgraph());
+  EXPECT_EQ(writer.SetCustomInputOutput(/*inputs=*/{0}, /*outputs=*/{2},
+                                        /*execution_plan=*/{0}),
+            kTfLiteOk);
+  writer.Write(test_file);
+
+  // Read model and test.
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  CHECK_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+}
+
+TEST_P(SingleSubgraphTest, PerTensorQuantizedModelTest) {
+  Interpreter interpreter;
+  interpreter.AddTensors(3);
+  interpreter.SetTensorParametersReadWrite(
+      0, kTfLiteUInt8, "a", {3}, TfLiteQuantizationParams({1 / 256., 128}));
+  interpreter.SetTensorParametersReadWrite(
+      1, kTfLiteUInt8, "b", {3}, TfLiteQuantizationParams({1 / 256., 128}));
+  interpreter.SetTensorParametersReadWrite(
+      2, kTfLiteUInt8, "c", {3}, TfLiteQuantizationParams({1 / 256., 128}));
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({2});
+  const char* initial_data = "";
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+  interpreter.AddNodeWithParameters({0, 1}, {2}, initial_data, 0,
+                                    reinterpret_cast<void*>(builtin_data), reg);
+
+  const std::string test_file = CreateFilePath("test_uint8.tflite");
+  WriteToFile(&interpreter, test_file, GetParam());
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  CHECK_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+}
+
+TEST_P(SingleSubgraphTest, OpVersioningTest) {
+  Interpreter interpreter;
+  interpreter.AddTensors(3);
+  interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "a", {1, 4},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadWrite(1, kTfLiteInt32, "b", {2},
+                                           TfLiteQuantization());
+  interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "c", {4, 4},
+                                           TfLiteQuantization());
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({2});
+
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  const TfLiteRegistration* reg =
+      resolver.FindOp(BuiltinOperator_BROADCAST_TO, 2);
+  interpreter.AddNodeWithParameters(/*inputs=*/{0, 1}, /*outputs=*/{2},
+                                    /*init_data=*/nullptr, /*init_data_size=*/0,
+                                    /*builtin_data=*/nullptr, reg);
+
+  const std::string test_file = CreateFilePath("test_float.tflite");
+  WriteToFile(&interpreter, test_file, GetParam());
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  CHECK_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+
+  ASSERT_EQ(new_interpreter->nodes_size(), 1);
+  TfLiteRegistration output_reg =
+      new_interpreter->node_and_registration(0)->second;
+  ASSERT_EQ(output_reg.builtin_code, BuiltinOperator_BROADCAST_TO);
+  CHECK_EQ(output_reg.version, 2);
+}
+
+TEST_P(SingleSubgraphTest, DynamicShapeTest) {
+  // Build a model with a single Add op.
+  Interpreter interpreter;
+
+  interpreter.AddTensors(3);
+  std::vector<int> dims = {1, 3};
+  std::vector<int> dims_signature = {-1, 3};
+  interpreter.SetTensorParametersReadWrite(
+      0, kTfLiteFloat32, "a", dims, TfLiteQuantizationParams{1.0, 0},
+      /*is_variable=*/false, &dims_signature);
+  interpreter.SetTensorParametersReadWrite(
+      1, kTfLiteFloat32, "b", dims, TfLiteQuantizationParams{1.0, 0},
+      /*is_variable=*/false, &dims_signature);
+  interpreter.SetTensorParametersReadWrite(
+      2, kTfLiteFloat32, "c", dims, TfLiteQuantizationParams{1.0, 0},
+      /*is_variable=*/false, &dims_signature);
+
+  interpreter.SetInputs({0, 1});
+  interpreter.SetOutputs({2});
+  const char* initial_data = "";
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  TfLiteAddParams* builtin_data =
+      reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
+  builtin_data->activation = kTfLiteActNone;
+  builtin_data->pot_scale_int16 = false;
+  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+  interpreter.AddNodeWithParameters({0, 1}, {2}, initial_data, 0,
+                                    reinterpret_cast<void*>(builtin_data), reg);
+
+  // Export interpreter and import back.
+  const std::string test_file = CreateFilePath("test_dynamic_shape.tflite");
+  WriteToFile(&interpreter, test_file, GetParam());
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  CHECK_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+
+  // Check shape signature in new interpreter.
+  TfLiteTensor* tensor0 = new_interpreter->tensor(0);
+  CHECK_NOTNULL(tensor0->dims_signature);
+  TfLiteIntArrayView shape_view(tensor0->dims_signature);
+  CHECK_EQ(shape_view.size(), 2);
+  CHECK_EQ(shape_view[0], -1);
+}
+
+INSTANTIATE_TEST_SUITE_P(Writer, SingleSubgraphTest, ::testing::Bool());
+
+struct ReshapeTestPattern {
+  int num_inputs;
+  bool is_param_valid;
+  bool has_buggy_non_flatten_shape;
+};
+
+class ReshapeLayerTest : public ::testing::TestWithParam<ReshapeTestPattern> {};
+
+TEST_P(ReshapeLayerTest, ReshapeLayerTest) {
+  const auto param = GetParam();
+  Interpreter interpreter;
+  const int total_tensors = param.num_inputs + 1;
+  interpreter.AddTensors(total_tensors);
+  int output_shape[] = {1, 2, 3};
+  interpreter.SetTensorParametersReadWrite(/*tensor_index=*/0, kTfLiteFloat32,
+                                           /*name=*/"a", /*dims=*/{6},
+                                           TfLiteQuantization());
+  ASSERT_LE(param.num_inputs, 2);
+  if (param.num_inputs == 2) {
+    // Some TOCO generated models have buggy shape arguments, which are required
+    // to be flatten, for example, dims={3, 1} instead of dims={3}.
+    if (param.has_buggy_non_flatten_shape) {
+      interpreter.SetTensorParametersReadOnly(
+          /*tensor_index=*/1, kTfLiteInt32, /*name=*/"b", /*dims=*/{3, 1},
+          TfLiteQuantization(), reinterpret_cast<char*>(output_shape),
+          sizeof(output_shape));
+    } else {
+      interpreter.SetTensorParametersReadOnly(
+          /*tensor_index=*/1, kTfLiteInt32, /*name=*/"b", /*dims=*/{3},
+          TfLiteQuantization(), reinterpret_cast<char*>(output_shape),
+          sizeof(output_shape));
+    }
+  }
+  interpreter.SetTensorParametersReadWrite(/*tensor_index=*/total_tensors - 1,
+                                           kTfLiteFloat32, /*name=*/"c",
+                                           /*dims=*/{3}, TfLiteQuantization());
+
+  std::vector<int> input_tensors(param.num_inputs);
+  std::iota(input_tensors.begin(), input_tensors.end(), 0);
+
+  interpreter.SetInputs(input_tensors);
+  interpreter.SetOutputs({total_tensors - 1});
+  const char* initial_data = "";
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  TfLiteReshapeParams* builtin_data = reinterpret_cast<TfLiteReshapeParams*>(
+      malloc(sizeof(TfLiteReshapeParams)));
+  memset(builtin_data, 0, sizeof(TfLiteReshapeParams));
+  if (param.is_param_valid) {
+    builtin_data->num_dimensions = 3;
+    for (int dim = 0; dim < builtin_data->num_dimensions; ++dim) {
+      builtin_data->shape[dim] = output_shape[dim];
+    }
+  }
+  const TfLiteRegistration* reg = resolver.FindOp(BuiltinOperator_RESHAPE, 1);
+  interpreter.AddNodeWithParameters(input_tensors,
+                                    /*outputs=*/{total_tensors - 1},
+                                    initial_data, /*init_data_size=*/0,
+                                    reinterpret_cast<void*>(builtin_data), reg);
+
+  SubgraphWriter writer(&interpreter.primary_subgraph());
+  std::stringstream ss;
+  ss << CreateFilePath("test_reshape_") << param.num_inputs
+     << param.is_param_valid << ".tflite";
+  std::string filename = ss.str();
+  writer.Write(filename);
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(filename.c_str());
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+  ASSERT_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    Writer, ReshapeLayerTest,
+    ::testing::Values(ReshapeTestPattern{/*num_inputs=*/2,
+                                         /*is_param_valid=*/true,
+                                         /*has_buggy_non_flatten_shape=*/false},
+                      ReshapeTestPattern{/*num_inputs=*/2,
+                                         /*is_param_valid=*/false,
+                                         /*has_buggy_non_flatten_shape=*/false},
+                      ReshapeTestPattern{/*num_inputs=*/1,
+                                         /*is_param_valid=*/true,
+                                         /*has_buggy_non_flatten_shape=*/false},
+                      ReshapeTestPattern{/*num_inputs=*/2,
+                                         /*is_param_valid=*/true,
+                                         /*has_buggy_non_flatten_shape=*/true}),
+    [](const ::testing::TestParamInfo<ReshapeLayerTest::ParamType>& info) {
+      std::stringstream ss;
+      ss << "num_inputs_" << info.param.num_inputs << "_valid_param_"
+         << info.param.is_param_valid << "_buggy_shape_"
+         << info.param.has_buggy_non_flatten_shape;
+      std::string name = ss.str();
+      return name;
+    });
+
+class WhileTest : public subgraph_test_util::ControlFlowOpTest {
+ protected:
+  TfLiteCustomAllocation NewCustomAlloc(size_t num_bytes,
+                                        int required_alignment) {
+    // Extra memory to ensure alignment.
+    char* new_alloc = new char[num_bytes + required_alignment];
+    char* new_underlying_buffer_aligned_ptr = reinterpret_cast<char*>(
+        AlignTo(required_alignment, reinterpret_cast<intptr_t>(new_alloc)));
+    custom_alloc_buffers_.emplace_back(new_alloc);
+
+    return TfLiteCustomAllocation(
+        {new_underlying_buffer_aligned_ptr, num_bytes});
+  }
+
+  intptr_t AlignTo(size_t alignment, intptr_t offset) {
+    return offset % alignment == 0 ? offset
+                                   : offset + (alignment - offset % alignment);
+  }
+
+  std::vector<std::unique_ptr<char[]>> custom_alloc_buffers_;
+};
+
+// The test builds a model that produces the i-th number of
+// triangular number sequence: 1, 3, 6, 10, 15, 21, 28.
+TEST_F(WhileTest, TestTriangularNumberSequence) {
+  const int kSeqNumber = 4;
+  const int kExpectedValue = 15;
+
+  interpreter_ = std::make_unique<Interpreter>();
+  AddSubgraphs(2);
+  builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(1), kSeqNumber);
+  builder_->BuildAccumulateLoopBodySubgraph(interpreter_->subgraph(2));
+  builder_->BuildWhileSubgraph(&interpreter_->primary_subgraph());
+
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1});
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+  FillIntTensor(interpreter_->tensor(interpreter_->inputs()[0]), {1});
+
+  // Use custom allocation for second input, to ensure things work well for
+  // non-traditional allocation types.
+  auto alloc =
+      NewCustomAlloc(interpreter_->tensor(interpreter_->inputs()[1])->bytes,
+                     kDefaultTensorAlignment);
+  auto* input_data = reinterpret_cast<int*>(alloc.data);
+  input_data[0] = 1;
+  interpreter_->SetCustomAllocationForTensor(interpreter_->inputs()[1], alloc);
+
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+  TfLiteTensor* output1 = interpreter_->tensor(interpreter_->outputs()[0]);
+  CheckIntTensor(output1, {1}, {kSeqNumber + 1});
+  TfLiteTensor* output2 = interpreter_->tensor(interpreter_->outputs()[1]);
+  CheckIntTensor(output2, {1}, {kExpectedValue});
+
+  // Now serialize & deserialize model into a new Interpreter.
+  ModelWriter writer(interpreter_.get());
+  const std::string test_file = CreateFilePath("test_while.tflite");
+  writer.Write(test_file);
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+
+  // Check deserialized model.
+  new_interpreter->ResizeInputTensor(new_interpreter->inputs()[0], {1});
+  new_interpreter->ResizeInputTensor(new_interpreter->inputs()[1], {1});
+  ASSERT_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+  FillIntTensor(new_interpreter->tensor(new_interpreter->inputs()[0]), {1});
+  FillIntTensor(new_interpreter->tensor(new_interpreter->inputs()[1]), {1});
+  ASSERT_EQ(new_interpreter->Invoke(), kTfLiteOk);
+  output1 = new_interpreter->tensor(new_interpreter->outputs()[0]);
+  CheckIntTensor(output1, {1}, {kSeqNumber + 1});
+  output2 = new_interpreter->tensor(new_interpreter->outputs()[1]);
+  CheckIntTensor(output2, {1}, {kExpectedValue});
+}
+
+// Verifies the ModelWriters constructing from  an interpreter or subgraphs
+// produce the same results.
+TEST_F(WhileTest, TestModelWriterFromSubgraphs) {
+  const int kSeqNumber = 4;
+  const int kExpectedValue = 15;
+
+  interpreter_ = std::make_unique<Interpreter>();
+  AddSubgraphs(2);
+  builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(1), kSeqNumber);
+  builder_->BuildAccumulateLoopBodySubgraph(interpreter_->subgraph(2));
+  builder_->BuildWhileSubgraph(&interpreter_->primary_subgraph());
+
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1});
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+  FillIntTensor(interpreter_->tensor(interpreter_->inputs()[0]), {1});
+
+  // Use custom allocation for second input, to ensure things work well for
+  // non-traditional allocation types.
+  auto alloc =
+      NewCustomAlloc(interpreter_->tensor(interpreter_->inputs()[1])->bytes,
+                     kDefaultTensorAlignment);
+  auto* input_data = reinterpret_cast<int*>(alloc.data);
+  input_data[0] = 1;
+  interpreter_->SetCustomAllocationForTensor(interpreter_->inputs()[1], alloc);
+
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+  TfLiteTensor* output1 = interpreter_->tensor(interpreter_->outputs()[0]);
+  CheckIntTensor(output1, {1}, {kSeqNumber + 1});
+  TfLiteTensor* output2 = interpreter_->tensor(interpreter_->outputs()[1]);
+  CheckIntTensor(output2, {1}, {kExpectedValue});
+
+  // Serializes the model using the interpreter.
+  ModelWriter writer_1(interpreter_.get());
+  const std::string test_file_1 = CreateFilePath("test_while_1.tflite");
+  writer_1.Write(test_file_1);
+
+  // Serializes the model using subgraphs.
+  std::vector<Subgraph*> subgraphs;
+  for (int i = 0; i < interpreter_->subgraphs_size(); ++i) {
+    subgraphs.push_back(interpreter_->subgraph(i));
+  }
+  ModelWriter writer_2(subgraphs);
+  const std::string test_file_2 = CreateFilePath("test_while_2.tflite");
+  writer_2.Write(test_file_2);
+
+  // The results from both ModelWriters should be the same.
+  std::ifstream file_ifs_1(test_file_1, std::ios::in);
+  std::ostringstream model_content_1;
+  model_content_1 << file_ifs_1.rdbuf();
+
+  std::ifstream file_ifs_2(test_file_2, std::ios::in);
+  std::ostringstream model_content_2;
+  model_content_2 << file_ifs_2.rdbuf();
+
+  EXPECT_FALSE(model_content_1.str().empty());
+  EXPECT_EQ(model_content_1.str(), model_content_2.str());
+}
+
+// The test builds a model with two while loops and only serializes one.
+// Loops compute the triangular numbers of different sizes.
+TEST_F(WhileTest, TestUpdateSubgraphIndices) {
+  const int kSeqNumber1 = 4;
+  const int kSeqNumber2 = 5;
+  const int kExpectedValue1 = 15;
+  const int kExpectedValue2 = 21;
+
+  interpreter_ = std::make_unique<Interpreter>();
+  AddSubgraphs(4);
+  builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(1), kSeqNumber1);
+  builder_->BuildAccumulateLoopBodySubgraph(interpreter_->subgraph(2));
+  builder_->BuildLessEqualCondSubgraph(interpreter_->subgraph(3), kSeqNumber2);
+  builder_->BuildAccumulateLoopBodySubgraph(interpreter_->subgraph(4));
+
+  // Build the primary subgraph.
+  // kInput1(0) --> +--------+ --> kUnused1(2)
+  //                | WHILE1 |
+  // kInput2(1) --> +--------+ --> kOutput1(4)
+
+  // kInput1(0) --> +--------+ --> kUnused2(3)
+  //                | WHILE2 |
+  // kInput2(1) --> +--------+ --> kOutput2(5)
+  Subgraph* primary_subgraph = &interpreter_->primary_subgraph();
+
+  const int kInput1 = 0;
+  const int kInput2 = 1;
+  const int kUnused1 = 2;
+  const int kUnused2 = 3;
+  const int kOutput1 = 4;
+  const int kOutput2 = 5;
+  const int kTensorCount = 6;
+
+  int first_new_tensor_index;
+  ASSERT_EQ(primary_subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
+            kTfLiteOk);
+  ASSERT_EQ(first_new_tensor_index, 0);
+  ASSERT_EQ(primary_subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
+  ASSERT_EQ(primary_subgraph->SetOutputs({kOutput1, kOutput2}), kTfLiteOk);
+  for (int i = 0; i < kTensorCount; ++i) {
+    ASSERT_EQ(primary_subgraph->SetTensorParametersReadWrite(
+                  i, kTfLiteInt32, "", 0, nullptr, {}, false),
+              kTfLiteOk);
+  }
+
+  // Register and add while ops.
+  auto* while_reg = ops::builtin::Register_WHILE();
+  while_reg->builtin_code = kTfLiteBuiltinWhile;
+
+  TfLiteWhileParams* params1 =
+      reinterpret_cast<TfLiteWhileParams*>(malloc(sizeof(TfLiteWhileParams)));
+  params1->cond_subgraph_index = 1;
+  params1->body_subgraph_index = 2;
+
+  TfLiteWhileParams* params2 =
+      reinterpret_cast<TfLiteWhileParams*>(malloc(sizeof(TfLiteWhileParams)));
+  params2->cond_subgraph_index = 3;
+  params2->body_subgraph_index = 4;
+
+  int while1_index, while2_index;
+  primary_subgraph->AddNodeWithParameters({kInput1, kInput2},
+                                          {kUnused1, kOutput1}, {}, nullptr, 0,
+                                          params1, while_reg, &while1_index);
+  primary_subgraph->AddNodeWithParameters({kInput1, kInput2},
+                                          {kUnused2, kOutput2}, {}, nullptr, 0,
+                                          params2, while_reg, &while2_index);
+
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1});
+  interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1});
+  ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+  FillIntTensor(interpreter_->tensor(interpreter_->inputs()[0]), {1});
+
+  // Use custom allocation for second input, to ensure things work well for
+  // non-traditional allocation types.
+  auto alloc =
+      NewCustomAlloc(interpreter_->tensor(interpreter_->inputs()[1])->bytes,
+                     kDefaultTensorAlignment);
+  auto* input_data = reinterpret_cast<int*>(alloc.data);
+  input_data[0] = 1;
+  interpreter_->SetCustomAllocationForTensor(interpreter_->inputs()[1], alloc);
+
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+  TfLiteTensor* output1 = interpreter_->tensor(interpreter_->outputs()[0]);
+  CheckIntTensor(output1, {1}, {kExpectedValue1});
+  TfLiteTensor* output2 = interpreter_->tensor(interpreter_->outputs()[1]);
+  CheckIntTensor(output2, {1}, {kExpectedValue2});
+
+  // Now serialize & deserialize model into a new Interpreter.
+  // Only serialize the second while op and related subgraphs.
+  ModelWriter writer({interpreter_->subgraph(0), interpreter_->subgraph(3),
+                      interpreter_->subgraph(4)});
+  writer.SetCustomInputOutput(/*subgraph_index=*/0, {kInput1, kInput2},
+                              {kOutput2}, {while2_index});
+  const std::string test_file = CreateFilePath("test_while.tflite");
+  writer.Write(test_file);
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromFile(test_file.c_str());
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  InterpreterBuilder builder(*model, resolver);
+  std::unique_ptr<Interpreter> new_interpreter;
+  builder(&new_interpreter);
+
+  // Check deserialized model.
+  new_interpreter->ResizeInputTensor(new_interpreter->inputs()[0], {1});
+  new_interpreter->ResizeInputTensor(new_interpreter->inputs()[1], {1});
+  ASSERT_EQ(new_interpreter->AllocateTensors(), kTfLiteOk);
+  FillIntTensor(new_interpreter->tensor(new_interpreter->inputs()[0]), {1});
+  FillIntTensor(new_interpreter->tensor(new_interpreter->inputs()[1]), {1});
+  ASSERT_EQ(new_interpreter->Invoke(), kTfLiteOk);
+  ASSERT_EQ(new_interpreter->outputs().size(), 1);
+  TfLiteTensor* output = new_interpreter->tensor(new_interpreter->outputs()[0]);
+  CheckIntTensor(output, {1}, {kExpectedValue2});
+}
+
+}  // namespace tflite

tensorflow/lite/tools/serialization/writer_test.cc

@@ -0,0 +1,60 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+// Loads the input tflite file into interpreter, serializes it back to a tflite
+// buffer, and then verifies that the generated output can be loaded back into
+// an interpreter and the model prepared (i.e., AllocateTensors returns ok).
+//
+// Usage:
+//   writer_test <input tflite>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <memory>
+
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/tools/serialization/writer_lib.h"
+
+int main(int argc, char* argv[]) {
+  if (argc != 2) {
+    fprintf(stderr, "Usage: %s input_file\n", argv[0]);
+    return 1;
+  }
+  std::unique_ptr<tflite::FlatBufferModel> model =
+      tflite::FlatBufferModel::BuildFromFile(argv[1]);
+  std::unique_ptr<tflite::Interpreter> interpreter;
+  tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates
+      builtin_op_resolver;
+  tflite::InterpreterBuilder(*model, builtin_op_resolver)(&interpreter);
+  tflite::ModelWriter writer(interpreter.get());
+  std::unique_ptr<uint8_t[]> output_buffer;
+  size_t output_buffer_size;
+  writer.GetBuffer(&output_buffer, &output_buffer_size);
+
+  // Verify the generated model.
+  std::unique_ptr<tflite::Interpreter> new_interpreter;
+  model = tflite::FlatBufferModel::BuildFromBuffer(
+      reinterpret_cast<char*>(output_buffer.get()), output_buffer_size);
+  tflite::InterpreterBuilder(*model, builtin_op_resolver)(&new_interpreter);
+  if (new_interpreter->AllocateTensors() != kTfLiteOk) {
+    fprintf(stderr, "AllocateTensors failed on the round-tripped model.\n");
+    return 1;
+  }
+  return 0;
+}