chore: import upstream snapshot with attribution

tensorflow/tools/proto_splitter/BUILD

@@ -0,0 +1,253 @@
+# Description:
+#   Utilities for splitting and joining large protos > 2GB.
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_oss",
+    "py_test",
+    "tf_cc_test",
+)
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "__subpackages__",
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+tf_proto_library(
+    name = "versions_proto",
+    srcs = ["versions.proto"],
+    make_default_target_header_only = True,
+)
+
+tf_proto_library(
+    name = "chunk_proto",
+    srcs = ["chunk.proto"],
+    make_default_target_header_only = True,
+    protodeps = [":versions_proto"],
+)
+
+# OSS-only: Required to link the impl files.
+cc_library(
+    name = "protos_impl",
+    deps = if_oss([
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc_impl",
+        "//tensorflow/tools/proto_splitter:versions_proto_cc_impl",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc_impl",
+    ]),
+)
+
+# copybara:uncomment_begin(google-only)
+#
+# py_proto_library(
+#     name = "versions_proto_py_pb2",
+#     deps = [
+#         ":versions_proto",
+#     ],
+# )
+#
+# py_proto_library(
+#     name = "chunk_proto_py_pb2",
+#     deps = [
+#         ":chunk_proto",
+#     ],
+# )
+#
+# copybara:uncomment_end
+
+pytype_strict_library(
+    name = "split",
+    srcs = ["split.py"],
+    deps = [
+        ":chunk_proto_py",
+        ":util",
+        ":version",
+        ":versions_proto_py",
+        "//tensorflow/python/lib/io:file_io",
+        "@absl_py//absl/logging",
+        "@riegeli//python/riegeli",
+    ],
+)
+
+pytype_strict_library(
+    name = "version",
+    srcs = ["version.py"],
+)
+
+py_test(
+    name = "split_test",
+    srcs = ["split_test.py"],
+    strict_deps = True,
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":chunk_proto_py",
+        ":split",
+        ":versions_proto_py",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_py",
+        "@absl_py//absl/testing:parameterized",
+        "@riegeli//python/riegeli",
+    ],
+)
+
+pytype_strict_library(
+    name = "constants",
+    srcs = ["constants.py"],
+    visibility = ["//visibility:public"],
+)
+
+pytype_strict_library(
+    name = "split_graph_def",
+    srcs = ["split_graph_def.py"],
+    deps = [
+        ":constants",
+        ":split",
+        ":util",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/framework:tensor_util",
+    ],
+)
+
+py_test(
+    name = "split_graph_def_test",
+    srcs = ["split_graph_def_test.py"],
+    strict_deps = True,
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":chunk_proto_py",
+        ":constants",
+        ":split_graph_def",
+        ":util",
+        #internal proto upb dep
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/tools/proto_splitter/python:test_util",
+    ],
+)
+
+pytype_strict_library(
+    name = "util",
+    srcs = ["util.py"],
+    deps = [":chunk_proto_py"],
+)
+
+py_test(
+    name = "util_test",
+    srcs = ["util_test.py"],
+    strict_deps = True,
+    deps = [
+        ":util",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_py",
+    ],
+)
+
+cc_library(
+    name = "merge_impl",
+    srcs = [
+        "merge.cc",
+        "merge.h",
+    ],
+    deps = [
+        ":chunk_proto_cc",
+        "//tensorflow/core:lib",
+        "//tensorflow/tools/proto_splitter/cc:util",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@riegeli//riegeli/base:object",
+        "@riegeli//riegeli/bytes:fd_reader",
+        "@riegeli//riegeli/bytes:string_reader",
+        "@riegeli//riegeli/records:record_reader",
+        "@tsl//tsl/platform:protobuf",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+cc_library(
+    name = "merge",
+    hdrs = ["merge.h"],
+    visibility = [
+        "//tensorflow:internal",
+        "//tensorflow/cc/experimental/tf2:__subpackages__",
+        "//tensorflow/cc/saved_model/image_format:__subpackages__",
+    ],
+    deps = [
+        ":chunk_proto_cc",
+        ":merge_impl",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@riegeli//riegeli/bytes:fd_reader",
+        "@riegeli//riegeli/records:record_reader",
+        "@tsl//tsl/platform:protobuf",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+tf_cc_test(
+    name = "merge_test",
+    srcs = ["merge_test.cc"],
+    data = [
+        "//tensorflow/cc/saved_model:testdata/chunked_saved_model/chunked_model/saved_model.cpb",
+        "//tensorflow/cc/saved_model:testdata/chunked_saved_model/chunked_model/saved_model.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:bf-split-tree.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:bf-split-tree.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:df-split-tree.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:df-split-tree.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:many-field.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:many-field.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:split-large-constant.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:split-large-constant.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:split-large-nodes.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:split-large-nodes.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:split-lots-nodes.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:split-lots-nodes.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:split-standard.cpb",
+        "//tensorflow/tools/proto_splitter/testdata:split-standard.pbtxt",
+        "//tensorflow/tools/proto_splitter/testdata:split-tree.pbtxt",
+    ],
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":chunk_proto_cc",
+        ":merge",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/tools/proto_splitter/cc:test_util",
+        "//tensorflow/tools/proto_splitter/cc:util",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:env",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/tsl/platform:env",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)

tensorflow/tools/proto_splitter/README.md

@@ -0,0 +1,121 @@
+# Proto Splitter
+
+Utilities for splitting large protos.
+
+For a more detailed overview of the library, see our [in-depth guide](g3doc/in-depth-guide.md).
+
+## The Python `Splitter` class
+
+Users can apply the Splitter implementations by calling:
+
+```python
+splitter = MySplitterClass(proto)
+
+# Export the chunks to a file.
+splitter.write(file_prefix)
+
+# Access the chunks created in the splitter.
+chunks, chunked_message = splitter.split()
+```
+
+### Composable Riegeli splitter
+
+The `split.py` class provides a `ComposableSplitter` class that is implemented
+to write to the Riegeli format, and allows combinable implementations of
+different message splitters.
+
+Recommended steps to subclass `ComposableSplitter`:
+
+1.  (required) Override `build_chunks()`. This method sets the values of
+    `self._chunks` and `self._chunked_message` based on the user-passed proto.
+2.  Update `version_def`. This is important to ensure that users are able to
+    apply the Merger to the chunked proto, or get understandable version errors.
+3.  If `__init__` is overridden: call `super().__init__(proto, **kwargs)`. This
+    is optional but highly recommended since it sets up basic attributes that
+    may be needed by other splitters.
+
+#### Example
+
+Consider the `SavedModel` protobuf, which is composed of many different messages
+(shown below). It contains two message types that can be at risk of running `>
+2GB`: `SavedObjectGraph` and `GraphDef`.
+
+We can write a `SavedModelSplitter` that contains logic for chunking both types,
+or re-use splitter that specifically work on each. Considering that `GraphDef`
+is used widely outside of `SavedModel`, the latter option is preferable.
+
+```proto
+message SavedModel {
+  ...
+  repeated MetaGraphDef meta_graphs = 2;
+}
+message MetaGraphDef {
+  ...
+  GraphDef graph_def = 2;
+  SavedObjectGraph object_graph_def = 7;
+}
+message GraphDef {
+  repeated NodeDef node = 1;
+  FunctionDefLibrary library = 2;
+  ...
+}
+```
+
+The SavedModel splitter implementation would look like:
+
+```python
+class SavedModelSplitter(ComposableSplitter):
+  def build_chunks(self):
+    ObjectGraphSplitter(
+      saved_model.meta_graphs[0].object_graph_def,
+      parent_splitter=self,
+      fields_in_parent=["meta_graphs", 0, "object_graph_def"]
+    ).build_chunks()
+
+    GraphDefSplitter(
+      saved_model.meta_graphs[0].graph_def,
+      parent_splitter=self,
+      fields_in_parent=["meta_graphs", 0, "graph_def"],
+    ).build_chunks()
+
+# See the results:
+A.split()  # [...chunks from B, ...chunks from C]
+```
+
+When B.split() and C.split() are called, chunks are added to A's chunk list, and
+A's ChunkedMessage proto is updated directly.
+
+## The C++ `Merger` class
+
+Once the proto has been split and written to disk using the aforementioned
+`Splitter` class, it can be merged back into its original form using these
+methods:
+
+```c++
+absl::Status Merger::Merge(
+  const std::vector<std::unique_ptr<tsl::protobuf::Message>>& chunks,
+  const ::proto_splitter::ChunkedMessage& chunked_message,
+  tsl::protobuf::Message* merged);
+
+absl::Status Merger::Read(std::string prefix, tsl::protobuf::Message* merged);
+```
+
+`Merger::Merge` requires the user to already have a collection of chunks in
+memory, while `Merger::Read` merges a chunked proto directly from disk. The
+methods can be used like so:
+
+```c++
+// Merge
+std::vector<std::unique_ptr<tsl::protobuf::Message>> my_chunks = GetMyChunks();
+::proto_splitter::ChunkedMessage chunked_message = GetMyChunkedMessage();
+my_project::MyProto my_proto;
+Merger::Merge(my_chunks, chunked_message, &my_proto);
+
+// Read
+my_project::MyOtherProto my_other_proto;
+Merger::Read("path/to/saved_model", &my_other_proto);
+```
+
+##### In-Depth Guide
+
+Looking for a more detailed overview of the library? See our [in-depth guide](g3doc/in-depth-guide.md).

tensorflow/tools/proto_splitter/cc/BUILD

@@ -0,0 +1,335 @@
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_oss",
+    "tf_cc_test",
+)
+
+# Description:
+#   Utilities for splitting and joining large protos > 2GB.
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//learning/brain/tfrt/aot:__subpackages__",
+        "//learning/brain/tfrt/tf_tpu:__subpackages__",
+        "//tensorflow/cc/saved_model/image_format:__subpackages__",
+        "//tensorflow/tools/proto_splitter:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "split",
+    hdrs = ["split.h"],
+    deps = [
+        ":util",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "//tensorflow/tools/proto_splitter:versions_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@tsl//tsl/platform:protobuf",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+cc_library(
+    name = "composable_splitter_base",
+    srcs = ["composable_splitter_base.cc"],
+    hdrs = ["composable_splitter_base.h"],
+    deps = [
+        ":max_size",
+        ":split",
+        ":util",
+        "//tensorflow/core:lib",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:cord",
+        "@riegeli//riegeli/base:initializer",
+        "@riegeli//riegeli/bytes:cord_writer",
+        "@riegeli//riegeli/bytes:fd_writer",
+        "@riegeli//riegeli/bytes:string_writer",
+        "@riegeli//riegeli/records:record_writer",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:protobuf",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+cc_library(
+    name = "composable_splitter",
+    hdrs = ["composable_splitter.h"],
+    deps = [
+        ":composable_splitter_base",
+        ":util",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "@tsl//tsl/platform:protobuf",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+tf_cc_test(
+    name = "composable_splitter_test",
+    srcs = ["composable_splitter_test.cc"],
+    data = [],
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":composable_splitter",
+        ":test_util",
+        ":util",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:cord",
+        "@com_google_googletest//:gtest_main",
+        "@riegeli//riegeli/base:initializer",
+        "@riegeli//riegeli/bytes:cord_reader",
+        "@riegeli//riegeli/bytes:fd_reader",
+        "@riegeli//riegeli/bytes:string_reader",
+        "@riegeli//riegeli/records:record_reader",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:status_matchers",
+        "@xla//xla/tsl/lib/core:status_test_util",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+cc_library(
+    name = "util",
+    srcs = ["util.cc"],
+    hdrs = ["util.h"],
+    visibility = [
+        "//learning/brain/tfrt/aot:__subpackages__",
+        "//tensorflow/cc/saved_model:__subpackages__",
+        "//tensorflow/cc/saved_model/image_format:__subpackages__",
+        "//tensorflow/tools/proto_splitter:__subpackages__",
+    ],
+    deps = [
+        "//tensorflow/core/platform",
+        "//tensorflow/core/platform:env",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:cord",
+        "@com_google_absl//absl/strings:str_format",
+        "@riegeli//riegeli/base:initializer",
+        "@riegeli//riegeli/base:types",
+        "@riegeli//riegeli/bytes:fd_reader",
+        "@riegeli//riegeli/bytes:string_reader",
+        "@riegeli//riegeli/records:record_reader",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:statusor",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+tf_cc_test(
+    name = "util_test",
+    srcs = ["util_test.cc"],
+    data = [
+        "//tensorflow/tools/proto_splitter/testdata:split-standard.cpb",
+    ],
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":test_util",
+        ":util",
+        "//tensorflow/core/platform:macros",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/core/platform:test",
+        "//tensorflow/tools/proto_splitter:chunk_proto_cc",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:status",
+        "@tsl//tsl/platform:status_matchers",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/tsl/lib/core:status_test_util",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+cc_library(
+    name = "graph_def_splitter",
+    srcs = ["graph_def_splitter.cc"],
+    hdrs = ["graph_def_splitter.h"],
+    deps = [
+        ":composable_splitter",
+        ":composable_splitter_base",
+        ":large_node_splitter",
+        ":max_size",
+        ":repeated_field_splitter",
+        ":size_splitter",
+        ":util",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "size_splitter",
+    hdrs = ["size_splitter.h"],
+    deps = [
+        ":composable_splitter",
+        ":util",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:optional",
+    ],
+)
+
+cc_library(
+    name = "repeated_field_splitter",
+    hdrs = ["repeated_field_splitter.h"],
+    deps = [
+        ":composable_splitter",
+        ":max_size",
+        ":size_splitter",
+        ":util",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@tsl//tsl/platform:protobuf",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "graph_def_splitter_test",
+    srcs = ["graph_def_splitter_test.cc"],
+    data = [
+        "//tensorflow/tools/proto_splitter/testdata:function-large-nodes.pb",
+        "//tensorflow/tools/proto_splitter/testdata:function-lots-of-nodes.pb",
+        "//tensorflow/tools/proto_splitter/testdata:graph-def-and-function.pb",
+        "//tensorflow/tools/proto_splitter/testdata:split-large-constant.pb",
+        "//tensorflow/tools/proto_splitter/testdata:split-large-nodes.pb",
+        "//tensorflow/tools/proto_splitter/testdata:split-lots-nodes.pb",
+    ],
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":graph_def_splitter",
+        ":max_size",
+        ":test_util",
+        ":util",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
+        "@com_google_absl//absl/strings:cord",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:protobuf",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+cc_library(
+    name = "max_size",
+    srcs = ["max_size.cc"],
+    hdrs = ["max_size.h"],
+    deps = ["@com_google_absl//absl/synchronization"],
+)
+
+cc_library(
+    name = "saved_model_splitter",
+    srcs = ["saved_model_splitter.cc"],
+    hdrs = ["saved_model_splitter.h"],
+    deps = [
+        ":composable_splitter",
+        ":graph_def_splitter",
+        ":large_node_splitter",
+        ":max_size",
+        ":util",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:protobuf",
+    ],
+)
+
+tf_cc_test(
+    name = "saved_model_splitter_test",
+    srcs = ["saved_model_splitter_test.cc"],
+    data = [
+        "//tensorflow/cc/saved_model:saved_model_test_files",
+    ],
+    tags = ["no_windows"],  # b/291001524
+    deps = [
+        ":max_size",
+        ":saved_model_splitter",
+        ":util",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:protobuf",
+    ] + if_oss([
+        "//tensorflow/tools/proto_splitter:protos_impl",
+    ]),
+)
+
+cc_library(
+    name = "large_node_splitter",
+    hdrs = ["large_node_splitter.h"],
+    deps = [
+        ":composable_splitter",
+        ":composable_splitter_base",
+        ":max_size",
+        ":size_splitter",
+        ":util",
+        "@com_google_absl//absl/status:statusor",
+        "@tsl//tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "test_util",
+    testonly = True,
+    hdrs = ["test_util.h"],
+    deps = [
+        "//tensorflow/core:test",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:statusor",
+    ],
+)

tensorflow/tools/proto_splitter/cc/composable_splitter.h

@@ -0,0 +1,52 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_H_
+
+#include <vector>
+
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter_base.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// A Splitter that can be composed with other splitters.
+class ComposableSplitter : public ComposableSplitterBase {
+ public:
+  // Initializer.
+  explicit ComposableSplitter(tsl::protobuf::Message* message)
+      : ComposableSplitterBase(message), message_(message) {}
+
+  // Initialize a child ComposableSplitter.
+  explicit ComposableSplitter(tsl::protobuf::Message* message,
+                              ComposableSplitterBase* parent_splitter,
+                              std::vector<FieldType>* fields_in_parent)
+      : ComposableSplitterBase(message, parent_splitter, fields_in_parent),
+        message_(message) {}
+
+ protected:
+  tsl::protobuf::Message* message() { return message_; }
+
+ private:
+  tsl::protobuf::Message* message_;
+};
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_H_

tensorflow/tools/proto_splitter/cc/composable_splitter_base.cc

@@ -0,0 +1,321 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/composable_splitter_base.h"
+
+#include <unistd.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <deque>
+#include <iterator>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/cord.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "riegeli/base/maker.h"  // from @riegeli
+#include "riegeli/bytes/cord_writer.h"  // from @riegeli
+#include "riegeli/bytes/fd_writer.h"  // from @riegeli
+#include "riegeli/bytes/string_writer.h"  // from @riegeli
+#include "riegeli/records/record_writer.h"  // from @riegeli
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/tools/proto_splitter/cc/max_size.h"
+#include "tensorflow/tools/proto_splitter/cc/split.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/errors.h"
+
+#define IS_OSS true
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+using ::tensorflow::proto_splitter::ChunkMetadata;
+
+VersionDef ComposableSplitterBase::Version() {
+  VersionDef version;
+  version.set_splitter_version(1);
+  version.set_join_version(0);
+  return version;
+}
+
+void ComposableSplitterBase::SetInitialSize(size_t size) { size_ = size; }
+
+size_t ComposableSplitterBase::GetInitialSize() {
+  if (size_ == 0) {
+    size_ = message_->ByteSizeLong();
+  }
+  return size_;
+}
+
+absl::StatusOr<ChunkedProto> ComposableSplitterBase::Split() {
+  if (parent_splitter_ != nullptr) {
+    return absl::UnimplementedError(
+        "The `Split` function behavior for children ComposableSplitter has not "
+        "been defined. Please call `parent_splitter.Split()` instead.");
+  }
+  if (!built_) {
+    LOG(INFO) << "Splitting message '" << message_->GetDescriptor()->full_name()
+              << "' (size " << HumanReadableBytes(GetInitialSize()) << ") "
+              << "into chunks of size " << HumanReadableBytes(GetMaxSize());
+    uint64_t start_time = Env::Default()->NowMicros();
+    TF_RETURN_IF_ERROR(BuildChunks());
+    TF_RETURN_IF_ERROR(FixChunks());
+    uint64_t end_time = Env::Default()->NowMicros();
+
+    std::string chunk_msg;
+    if (chunked_message_.chunked_fields().empty()) {
+      chunk_msg = "No chunks were generated.";
+    } else {
+      chunk_msg = absl::StrCat(
+          "Generated ", chunked_message_.chunked_fields_size(), " chunks.");
+    }
+    LOG(INFO) << "Finished chunking '" << message_->GetDescriptor()->full_name()
+              << "', took " << HumanReadableDuration(end_time - start_time)
+              << ". " << chunk_msg;
+    built_ = true;
+  }
+  return (ChunkedProto){.chunks = &chunks_,
+                        .chunked_message = &chunked_message_};
+}
+
+static absl::Status WriteToRecordWriter(
+    riegeli::RecordWriterBase& writer, const std::vector<MessageBytes>& chunks,
+    ChunkedMessage& chunked_message,
+    const ::tensorflow::proto_splitter::VersionDef& version) {
+  // Export Riegeli / chunked file.
+  ChunkMetadata metadata;
+  *metadata.mutable_message() = chunked_message;
+  *metadata.mutable_version() = version;
+  auto* metadata_chunks = metadata.mutable_chunks();
+
+  for (const auto& chunk : chunks) {
+    auto* chunk_metadata = metadata_chunks->Add();
+    if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
+            chunk)) {
+      const auto& msg_chunk =
+          std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk);
+      LOG(INFO) << "Writing chunk of size " << msg_chunk->ByteSizeLong();
+      writer.WriteRecord(*msg_chunk);
+      chunk_metadata->set_size(msg_chunk->ByteSizeLong());
+      chunk_metadata->set_type(
+          ::tensorflow::proto_splitter::ChunkInfo::MESSAGE);
+    } else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) {
+      auto* msg_chunk = std::get<tsl::protobuf::Message*>(chunk);
+      writer.WriteRecord(*msg_chunk);
+      chunk_metadata->set_size(msg_chunk->ByteSizeLong());
+      chunk_metadata->set_type(
+          ::tensorflow::proto_splitter::ChunkInfo::MESSAGE);
+    } else {
+      const auto& str_chunk = std::get<std::string>(chunk);
+      writer.WriteRecord(str_chunk);
+      chunk_metadata->set_size(str_chunk.size());
+      chunk_metadata->set_type(::tensorflow::proto_splitter::ChunkInfo::BYTES);
+    }
+    chunk_metadata->set_offset(writer.LastPos().get().numeric());
+  }
+  writer.WriteRecord(metadata);
+  return absl::OkStatus();
+}
+
+absl::Status ComposableSplitterBase::CheckIfWriteImplemented() {
+  if (parent_splitter_ != nullptr) {
+    return absl::UnimplementedError(
+        "The `Write` function behavior for children ComposableSplitter has not "
+        "been defined. Please call `parent_splitter.Write()` instead.");
+  }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<std::string> ComposableSplitterBase::Write(
+    std::string file_prefix) {
+  TF_RETURN_IF_ERROR(CheckIfWriteImplemented());
+
+  auto split_results = Split();
+  if (!split_results.ok()) return split_results.status();
+
+  std::vector<MessageBytes>* chunks = split_results.value().chunks;
+  ChunkedMessage* chunked_message = split_results.value().chunked_message;
+
+  tsl::Env* env = tsl::Env::Default();
+  TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(
+      std::string{tensorflow::io::Dirname(file_prefix)}));
+
+  std::string output_path;
+  if (chunked_message->chunked_fields().empty()) {
+    // Export regular pb.
+    output_path = absl::StrCat(file_prefix, ".pb");
+    TF_RETURN_IF_ERROR(
+        tensorflow::WriteBinaryProto(env, output_path, *message_));
+  } else {
+    // Export Riegeli / chunked file.
+    output_path = absl::StrCat(file_prefix, ".cpb");
+    riegeli::RecordWriter writer(
+        riegeli::Maker<riegeli::FdWriter>(output_path));
+    if (!writer.is_open()) return writer.status();
+    TF_RETURN_IF_ERROR(
+        WriteToRecordWriter(writer, *chunks, *chunked_message, Version()));
+    if (!writer.Close()) return writer.status();
+  }
+  LOG(INFO) << "Splitter output written to " << output_path;
+  return output_path;
+}
+
+absl::StatusOr<std::tuple<std::string, bool>>
+ComposableSplitterBase::WriteToString() {
+  TF_RETURN_IF_ERROR(CheckIfWriteImplemented());
+
+  auto split_results = Split();
+  if (!split_results.ok()) return split_results.status();
+  std::vector<MessageBytes>* chunks = split_results.value().chunks;
+  ChunkedMessage* chunked_message = split_results.value().chunked_message;
+
+  std::string output;
+  if (chunked_message->chunked_fields().empty()) {
+    // Export regular pb.
+    if (!message_->SerializeToString(&output))
+      return absl::InvalidArgumentError("Serialization to string failed");
+    LOG(INFO) << "Splitter output written to string";
+    return std::make_tuple(output, false);
+  } else {
+    // Export Riegeli / chunked file.
+    riegeli::RecordWriter writer(
+        riegeli::Maker<riegeli::StringWriter>(&output));
+    if (!writer.is_open()) return writer.status();
+    TF_RETURN_IF_ERROR(
+        WriteToRecordWriter(writer, *chunks, *chunked_message, Version()));
+    if (!writer.Close()) return writer.status();
+    LOG(INFO) << "Splitter output written to string";
+    return std::make_tuple(output, true);
+  }
+}
+
+#if !IS_OSS
+absl::StatusOr<std::tuple<absl::Cord, bool>>
+ComposableSplitterBase::WriteToCord() {
+  TF_RETURN_IF_ERROR(CheckIfWriteImplemented());
+
+  auto split_results = Split();
+  if (!split_results.ok()) return split_results.status();
+  std::vector<MessageBytes>* chunks = split_results.value().chunks;
+  ChunkedMessage* chunked_message = split_results.value().chunked_message;
+
+  absl::Cord output;
+  if (chunked_message->chunked_fields().empty()) {
+    // Export regular pb.
+    if (!message_->SerializeToString(&output))
+      return absl::InvalidArgumentError("Serialization to absl::Cord failed");
+    LOG(INFO) << "Splitter output written to absl::Cord";
+    return std::make_tuple(output, false);
+  } else {
+    // Export Riegeli / chunked file.
+    riegeli::RecordWriter writer(riegeli::Maker<riegeli::CordWriter>(&output));
+    if (!writer.is_open()) return writer.status();
+    TF_RETURN_IF_ERROR(
+        WriteToRecordWriter(writer, *chunks, *chunked_message, Version()));
+    if (!writer.Close()) return writer.status();
+    LOG(INFO) << "Splitter output written to absl::Cord";
+    return std::make_tuple(output, true);
+  }
+}
+#endif
+
+absl::Status ComposableSplitterBase::SetMessageAsBaseChunk() {
+  if (!chunks_.empty()) {
+    return absl::FailedPreconditionError(
+        "Cannot set `message` as the base chunk since there are already "
+        "created chunks.");
+  }
+
+  chunks_.push_back(message_);
+  chunked_message_.set_chunk_index(0);
+  chunks_order_.push_back(0);
+  add_chunk_order_.push_back(0);
+  return absl::OkStatus();
+}
+
+absl::Status ComposableSplitterBase::AddChunk(
+    std::unique_ptr<MessageBytes> chunk, std::vector<FieldType>* fields,
+    int* index) {
+  if (parent_splitter_ != nullptr) {
+    std::vector<FieldType> all_fields(fields_in_parent_->begin(),
+                                      fields_in_parent_->end());
+    all_fields.insert(all_fields.end(), fields->begin(), fields->end());
+    return parent_splitter_->AddChunk(std::move(chunk), &all_fields, index);
+  }
+
+  auto new_chunk_index = chunks_.size();
+  auto new_field = chunked_message_.add_chunked_fields();
+  new_field->mutable_message()->set_chunk_index(new_chunk_index);
+  TF_RETURN_IF_ERROR(
+      AddFieldTag(*message_->GetDescriptor(), *fields, *new_field));
+
+  // Add chunk at the end or insert at the index position.
+  if (index == nullptr) {
+    chunks_.push_back(*chunk);
+    chunks_order_.push_back(new_chunk_index);
+  } else {
+    auto it = chunks_.begin();
+    std::advance(it, *index);
+    chunks_.insert(it, *chunk);
+
+    fix_chunk_order_ = true;
+    auto it2 = chunks_order_.begin();
+    std::advance(it2, *index);
+    chunks_order_.insert(it2, new_chunk_index);
+  }
+  add_chunk_order_.push_back(new_chunk_index);
+  return absl::OkStatus();
+}
+
+absl::Status ComposableSplitterBase::FixChunks() {
+  if (!fix_chunk_order_) return absl::OkStatus();
+
+  // Use `add_chunk_order_` and `chunks_order_` to update the chunk indices.
+  absl::flat_hash_map<int, int> chunk_indices;
+  for (int i = 0; i < chunks_order_.size(); ++i) {
+    chunk_indices[chunks_order_[i]] = i;
+  }
+
+  std::deque<ChunkedMessage*> to_fix = {&chunked_message_};
+  while (!to_fix.empty()) {
+    auto msg = to_fix.front();
+    to_fix.pop_front();
+    for (int i = 0; i < msg->chunked_fields_size(); ++i) {
+      to_fix.push_back(msg->mutable_chunked_fields(i)->mutable_message());
+    }
+
+    if (!msg->has_chunk_index()) continue;
+    int current_chunk_idx = msg->chunk_index();
+    int new_chunk_index = chunk_indices[add_chunk_order_[current_chunk_idx]];
+    msg->set_chunk_index(new_chunk_index);
+  }
+  fix_chunk_order_ = false;
+  return absl::OkStatus();
+}
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/cc/composable_splitter_base.h

@@ -0,0 +1,127 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_BASE_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_BASE_H_
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/cord.h"
+#include "tensorflow/tools/proto_splitter/cc/split.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tsl/platform/protobuf.h"
+
+#define IS_OSS true
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// Parts of the ComposableSplitter that are independent of the template (to
+// avoid linking issues).
+class ComposableSplitterBase : public Splitter {
+ public:
+  explicit ComposableSplitterBase(tsl::protobuf::Message* message)
+      : built_(false), message_(message), parent_splitter_(nullptr) {}
+
+  explicit ComposableSplitterBase(tsl::protobuf::Message* message,
+                                  ComposableSplitterBase* parent_splitter,
+                                  std::vector<FieldType>* fields_in_parent)
+      : built_(false),
+        message_(message),
+        parent_splitter_(parent_splitter),
+        fields_in_parent_(fields_in_parent) {}
+
+  // Splits a proto message.
+  // Returns a pair of (
+  //   chunks: List of messages/bytes,
+  //   ChunkedMessage: Metadata about the chunked fields.)
+  // If the message is not split, `chunks` should only contain the original
+  // message.
+  absl::StatusOr<ChunkedProto> Split() override;
+
+  // Serializes a proto to disk.
+  // The writer writes all chunks into a Riegeli file. The chunk metadata
+  // (ChunkMetadata) is written at the very end.
+  //   file_prefix: string prefix of the filepath. The writer will automatically
+  //     attach a `.pb` or `.cpb` (chunked pb) suffix depending on whether the
+  //     proto is split.
+  // Returns the full output path.
+  absl::StatusOr<std::string> Write(std::string file_prefix) override;
+  // The bool field record whether it's saved as a chunked protobuf (true) or
+  // regular protobuf (false).
+  absl::StatusOr<std::tuple<std::string, bool>> WriteToString();
+#if !IS_OSS
+  absl::StatusOr<std::tuple<absl::Cord, bool>> WriteToCord();
+#endif
+
+  VersionDef Version() override;
+
+  // Builds the Splitter object by generating chunks from the proto.
+  // Subclasses of `ComposableChunks` should only need to override this method.
+  // This method should be called once per Splitter to create the chunks.
+  // Users should call the methods `Split` or `Write` instead.
+  virtual absl::Status BuildChunks() = 0;
+
+  // Set or get the size of the initial user-provided proto.
+  void SetInitialSize(size_t size);
+  size_t GetInitialSize();
+
+ protected:
+  // Initializes `chunks_` and `chunked_message_`, with the user-provided
+  // message as the initial chunk. This step is optional.
+  absl::Status SetMessageAsBaseChunk();
+
+  // Adds a new chunk and updates the ChunkedMessage proto. If set, the index
+  // indicates where to insert the chunk.
+  absl::Status AddChunk(std::unique_ptr<MessageBytes> chunk,
+                        std::vector<FieldType>* fields, int* index = nullptr);
+
+  std::vector<FieldType>* fields_in_parent() { return fields_in_parent_; }
+
+ private:
+  // After BuildChunks(), the chunk_index of each nested ChunkedMessage is set
+  // to the length of the list when the chunk was added. This would be fine if
+  // the chunks were always added to the end of the list. However, this is not
+  // always the case the indices must be updated.
+  absl::Status FixChunks();
+  absl::Status CheckIfWriteImplemented();
+
+  bool built_;
+  tsl::protobuf::Message* message_;
+  std::vector<MessageBytes> chunks_;
+  ::tensorflow::proto_splitter::ChunkedMessage chunked_message_;
+  ComposableSplitterBase* parent_splitter_;
+  std::vector<FieldType>* fields_in_parent_;
+  size_t size_ = 0;
+
+  // Keep list of chunk keys in the order in which they were added to
+  // the list and the actual order of chunks. If the user inserts a
+  // chunk out of order, then these arrays will be used to fix the message chunk
+  // indices.
+  std::vector<int> add_chunk_order_;
+  std::vector<int> chunks_order_;
+  bool fix_chunk_order_ = false;
+};
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_BASE_H_

tensorflow/tools/proto_splitter/cc/composable_splitter_test.cc

@@ -0,0 +1,361 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/composable_splitter.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/cord.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "riegeli/base/maker.h"  // from @riegeli
+#include "riegeli/bytes/cord_reader.h"  // from @riegeli
+#include "riegeli/bytes/fd_reader.h"  // from @riegeli
+#include "riegeli/bytes/string_reader.h"  // from @riegeli
+#include "riegeli/records/record_reader.h"  // from @riegeli
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system_helper.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/tools/proto_splitter/cc/test_util.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+
+#define IS_OSS true
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+namespace {
+
+using ::tensorflow::proto_splitter::ChunkedMessage;
+using ::tensorflow::proto_splitter::ChunkMetadata;
+using ::tensorflow::proto_splitter_testdata::RepeatedRepeatedString;
+using ::tensorflow::proto_splitter_testdata::RepeatedString;
+using ::testing::HasSubstr;
+using ::testing::SizeIs;
+using tsl::testing::StatusIs;
+
+// Required in OSS to prevent string to bool conversion in FieldType variant.
+using namespace std::string_literals;  // NOLINT
+
+// Splits each string in a RepeatedString into separate chunks.
+class RepeatedStringSplitter : public ComposableSplitter {
+  friend class ComposableSplitter;
+
+ public:
+  using ComposableSplitter::ComposableSplitter;
+
+  absl::Status BuildChunks() override {
+    RepeatedString* repeated_string =
+        google::protobuf::DynamicCastMessage<RepeatedString>(message());
+    auto strings = repeated_string->strings();
+
+    if (strings.empty()) {
+      TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
+      return absl::OkStatus();
+    }
+    for (int i = 0; i < strings.size(); i++) {
+      auto s = std::make_unique<MessageBytes>(strings[i]);
+      std::vector<FieldType> fields = {"strings"s, i};
+      TF_RETURN_IF_ERROR(AddChunk(std::move(s), &fields));
+    }
+    return absl::OkStatus();
+  }
+};
+
+RepeatedString SetUpRepeatedString(std::vector<std::string> strings) {
+  RepeatedString message;
+  *message.mutable_strings() = {strings.begin(), strings.end()};
+  return message;
+}
+
+TEST(RepeatedStringSplitterTest, TestSplitChunks) {
+  std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
+  auto message = SetUpRepeatedString(strings);
+  RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
+  TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
+  std::vector<MessageBytes>* chunks = ret.chunks;
+  ASSERT_NE(chunks, nullptr);
+  ChunkedMessage* chunked_message = ret.chunked_message;
+  ASSERT_NE(chunked_message, nullptr);
+
+  for (int i = 0; i < chunks->size(); i++) {
+    MessageBytes chunk = (*chunks)[i];
+    EXPECT_THAT(chunk, ::testing::VariantWith<std::string>(strings[i]));
+  }
+  EXPECT_THAT(*chunked_message, EqualsProto(R"pb(chunked_fields {
+                                                   field_tag { field: 1 }
+                                                   field_tag { index: 0 }
+                                                   message { chunk_index: 0 }
+                                                 }
+                                                 chunked_fields {
+                                                   field_tag { field: 1 }
+                                                   field_tag { index: 1 }
+                                                   message { chunk_index: 1 }
+                                                 }
+                                                 chunked_fields {
+                                                   field_tag { field: 1 }
+                                                   field_tag { index: 2 }
+                                                   message { chunk_index: 2 }
+                                                 })pb"));
+
+  // Calling split again should return the same chunks/ChunkedMessage.
+  TF_ASSERT_OK_AND_ASSIGN(auto ret2, splitter.Split());
+  std::vector<MessageBytes>* chunks2 = ret2.chunks;
+  ChunkedMessage* chunked_message2 = ret2.chunked_message;
+  EXPECT_EQ(chunks2, chunks);
+  EXPECT_EQ(chunked_message2, chunked_message);
+}
+
+static void CheckChunks(riegeli::RecordReaderBase& reader,
+                        std::vector<std::string>& strings) {
+  ChunkMetadata chunk_metadata;
+  reader.Seek(reader.Size().value());
+  reader.SeekBack();
+  reader.ReadRecord(chunk_metadata);
+
+  auto& chunk_info = chunk_metadata.chunks();
+  EXPECT_EQ(chunk_info.size(), strings.size());
+  for (int i = 0; i < chunk_info.size(); i++) {
+    reader.Seek(chunk_info[i].offset());
+    absl::string_view chunk;
+    reader.ReadRecord(chunk);
+    EXPECT_EQ(strings[i], std::string(chunk));
+  }
+
+  EXPECT_THAT(chunk_metadata.message(),
+              EqualsProto(R"pb(chunked_fields {
+                                 field_tag { field: 1 }
+                                 field_tag { index: 0 }
+                                 message { chunk_index: 0 }
+                               }
+                               chunked_fields {
+                                 field_tag { field: 1 }
+                                 field_tag { index: 1 }
+                                 message { chunk_index: 1 }
+                               }
+                               chunked_fields {
+                                 field_tag { field: 1 }
+                                 field_tag { index: 2 }
+                                 message { chunk_index: 2 }
+                               })pb"));
+}
+
+TEST(RepeatedStringSplitterTest, TestWrite) {
+  std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
+  auto message = SetUpRepeatedString(strings);
+  RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
+
+  std::string output_prefix = tensorflow::io::GetTempFilename("");
+  TF_ASSERT_OK_AND_ASSIGN(std::string expected_file,
+                          splitter.Write(output_prefix));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto exists,
+                          internal::FileExists(Env::Default(), expected_file));
+  EXPECT_TRUE(exists);
+
+  // Look for the last chunk, which should contain a ChunkMetadata proto.
+  riegeli::RecordReader file_reader(
+      riegeli::Maker<riegeli::FdReader>(std::move(expected_file)));
+
+  CheckChunks(file_reader, strings);
+}
+
+TEST(RepeatedStringSplitterTest, TestWriteToString) {
+  std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
+  auto message = SetUpRepeatedString(strings);
+  RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
+  auto string_output_results = splitter.WriteToString();
+  TF_EXPECT_OK(string_output_results.status());
+  std::string string_output = std::get<0>(string_output_results.value());
+  bool is_chunked = std::get<1>(string_output_results.value());
+  EXPECT_TRUE(is_chunked);
+  // Look for the last chunk, which should contain a ChunkMetadata proto.
+  riegeli::RecordReader string_reader(
+      riegeli::Maker<riegeli::StringReader>(string_output));
+
+  CheckChunks(string_reader, strings);
+}
+
+#if !IS_OSS
+TEST(RepeatedStringSplitterTest, TestWriteToCord) {
+  std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
+  auto message = SetUpRepeatedString(strings);
+  RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
+  auto cord_output_results = splitter.WriteToCord();
+  TF_EXPECT_OK(cord_output_results.status());
+  absl::Cord cord_output = std::get<0>(cord_output_results.value());
+  bool is_chunked = std::get<1>(cord_output_results.value());
+  EXPECT_TRUE(is_chunked);
+  // Look for the last chunk, which should contain a ChunkMetadata proto.
+  riegeli::RecordReader cord_reader(
+      riegeli::Maker<riegeli::CordReader>(&cord_output));
+
+  CheckChunks(cord_reader, strings);
+}
+#endif
+
+TEST(RepeatedStringSplitterTest, TestNoSplit) {
+  RepeatedString message;  // No strings
+  RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
+  TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
+  std::vector<MessageBytes>* chunks = ret.chunks;
+  ASSERT_NE(chunks, nullptr);
+  ChunkedMessage* chunked_message = ret.chunked_message;
+  ASSERT_NE(chunked_message, nullptr);
+
+  EXPECT_THAT(*chunks, SizeIs(1));
+  EXPECT_THAT(*std::get<tsl::protobuf::Message*>((*chunks)[0]),
+              EqualsProto(""));
+  EXPECT_THAT(*chunked_message, EqualsProto(R"pb(chunk_index: 0)pb"));
+}
+
+// Splits each string in a RepeatedString into separate chunks.
+class RepeatedRepeatedStringSplitter : public ComposableSplitter {
+ public:
+  using ComposableSplitter::ComposableSplitter;
+
+  absl::Status BuildChunks() override {
+    TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
+    RepeatedRepeatedString* msg =
+        google::protobuf::DynamicCastMessage<RepeatedRepeatedString>(message());
+    auto repeated_strings = msg->rs();
+    for (int i = 0; i < repeated_strings.size(); i++) {
+      std::vector<FieldType> fields = {"rs"s, i};
+      auto splitter =
+          RepeatedStringSplitter(&repeated_strings[i], this, &fields);
+      TF_RETURN_IF_ERROR(splitter.BuildChunks());
+    }
+    return absl::OkStatus();
+  }
+};
+
+TEST(ComposableTest, RepeatedRepeatedStringTest) {
+  std::vector<std::string> strings1 = {"piece-1", "piece-2", "piece-3"};
+  auto rs1 = SetUpRepeatedString(strings1);
+  std::vector<std::string> strings2 = {"new-strings-1"};
+  auto rs2 = SetUpRepeatedString(strings2);
+  std::vector<std::string> strings3 = {"foo-1", "foo-2"};
+  auto rs3 = SetUpRepeatedString(strings3);
+
+  std::vector<RepeatedString> rs = {rs1, rs2, rs3};
+
+  RepeatedRepeatedString message;
+  message.mutable_rs()->Add(rs.begin(), rs.end());
+
+  RepeatedRepeatedStringSplitter splitter =
+      RepeatedRepeatedStringSplitter(&message);
+  TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
+  std::vector<MessageBytes>* chunks = ret.chunks;
+  ASSERT_NE(chunks, nullptr);
+  ChunkedMessage* chunked_message = ret.chunked_message;
+  ASSERT_NE(chunked_message, nullptr);
+
+  std::vector<std::string> expected_chunks = {
+      "piece-1", "piece-2", "piece-3", "new-strings-1", "foo-1", "foo-2"};
+
+  // RepeatedRepeatedStringSplitter sets the first chunk as the user-provided
+  // message, so the expected size is 7.
+  EXPECT_THAT(*chunks, SizeIs(7));
+  EXPECT_THAT(*std::get<tsl::protobuf::Message*>((*chunks)[0]),
+              EqualsProto(message));
+
+  for (int i = 1; i < chunks->size(); i++) {
+    MessageBytes chunk = (*chunks)[i];
+    EXPECT_THAT(chunk,
+                ::testing::VariantWith<std::string>(expected_chunks[i - 1]));
+  }
+
+  // message.rs[2].strings[0] (value = "foo-1") should be the chunk at index 5.
+  EXPECT_THAT(chunked_message->chunked_fields()[4],
+              EqualsProto(R"pb(field_tag { field: 2 }
+                               field_tag { index: 2 }
+                               field_tag { field: 1 }
+                               field_tag { index: 0 }
+                               message { chunk_index: 5 })pb"));
+}
+
+TEST(ComposableTest, ChildSplitterTest) {
+  std::vector<std::string> strings1 = {"piece-1", "piece-2", "piece-3"};
+  auto message1 = SetUpRepeatedString(strings1);
+  RepeatedStringSplitter splitter(&message1);
+  std::vector<FieldType> fields = {};
+
+  std::vector<std::string> strings2 = {"s1", "s2"};
+  auto message2 = SetUpRepeatedString(strings2);
+  RepeatedStringSplitter child(&message2, &splitter, &fields);
+
+  TF_EXPECT_OK(child.BuildChunks());
+  TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
+  std::vector<MessageBytes>* chunks = ret.chunks;
+  ASSERT_NE(chunks, nullptr);
+  EXPECT_THAT(*chunks, SizeIs(5));  // Total 5 chunks should be generated.
+}
+
+TEST(ComposableTest, ChildSplitterUnimplementedTest) {
+  RepeatedString message;
+  RepeatedStringSplitter splitter(&message);
+  std::vector<FieldType> fields = {};
+  RepeatedStringSplitter child(&message, &splitter, &fields);
+
+  EXPECT_THAT(child.Split(),
+              absl_testing::StatusIs(absl::StatusCode::kUnimplemented,
+                                     HasSubstr("`Split` function behavior")));
+  EXPECT_THAT(child.Write("str"),
+              absl_testing::StatusIs(absl::StatusCode::kUnimplemented,
+                                     HasSubstr("`Write` function behavior")));
+}
+
+class NoOpSplitter : public ComposableSplitter {
+ public:
+  using ComposableSplitter::ComposableSplitter;
+
+  absl::Status BuildChunks() override { return absl::OkStatus(); }
+};
+
+TEST(NoOpSplitterTest, TestWrite) {
+  std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
+  auto message = SetUpRepeatedString(strings);
+  NoOpSplitter splitter(&message);
+
+  std::string output_prefix = tensorflow::io::GetTempFilename("");
+  TF_ASSERT_OK_AND_ASSIGN(std::string expected_file,
+                          splitter.Write(output_prefix));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto exists,
+                          internal::FileExists(Env::Default(), expected_file));
+  EXPECT_TRUE(exists);
+
+  RepeatedString read_message;
+  auto status = tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                            expected_file, &read_message);
+
+  EXPECT_THAT(read_message, EqualsProto(message));
+}
+
+}  // namespace
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/cc/graph_def_splitter.cc

@@ -0,0 +1,241 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/graph_def_splitter.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter_base.h"
+#include "tensorflow/tools/proto_splitter/cc/large_node_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/max_size.h"
+#include "tensorflow/tools/proto_splitter/cc/repeated_field_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/size_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow::tools::proto_splitter {
+
+namespace {
+
+// Required in OSS to prevent string to bool conversion in FieldType variant.
+using namespace std::string_literals;  // NOLINT
+
+// Extracts the value of large constants into a new chunk.
+// Only works on NodeDef messages.
+class ConstantSplitter : public SizeSplitter {
+ public:
+  using SizeSplitter::SizeSplitter;
+  absl::StatusOr<int> BuildChunksReturnSize() override {
+    NodeDef* node = google::protobuf::DynamicCastMessage<NodeDef>(message());
+
+    std::vector<FieldType> tensor_field = {"attr"s, "value"s, "tensor"s};
+    std::vector<FieldType> content_field = {"attr"s, "value"s, "tensor"s,
+                                            "tensor_content"s};
+
+    TF_ASSIGN_OR_RETURN(auto ret, GetMutableField(node, tensor_field));
+    auto tensor_msg =
+        ret.parent->GetReflection()->MutableMessage(ret.parent, ret.field);
+    TensorProto* tensor_proto =
+        google::protobuf::DynamicCastMessage<TensorProto>(tensor_msg);
+
+    int size_diff;
+
+    if (tensor_proto->tensor_content().empty()) {
+      Tensor t;
+      if (!t.FromProto(*tensor_proto)) {
+        return absl::InvalidArgumentError(
+            "Invalid Const NodeDef.attr[\"value\"].tensor value.");
+      }
+
+      TensorProto container;
+      t.AsProtoTensorContent(&container);
+      size_diff = container.tensor_content().size();
+      auto x = std::make_unique<std::string>(
+          std::move(*container.mutable_tensor_content()));
+      auto y = std::make_unique<MessageBytes>(std::move(*x));
+      TF_RETURN_IF_ERROR(AddChunk(std::move(y), &content_field));
+    } else {
+      size_diff = tensor_proto->tensor_content().size();
+      auto x = std::make_unique<std::string>(
+          std::move(*tensor_proto->mutable_tensor_content()));
+      auto y = std::make_unique<MessageBytes>(std::move(*x));
+      TF_RETURN_IF_ERROR(AddChunk(std::move(y), &content_field));
+    }
+
+    // Keep the TensorProto's dtype, tensor_shape, and version_number fields,
+    // but clear the raw tensor content / "xxx_val" attributes.
+    auto dtype = tensor_proto->dtype();
+    auto tensor_shape = tensor_proto->tensor_shape();
+    auto version_number = tensor_proto->version_number();
+    tensor_proto->Clear();
+    tensor_proto->set_dtype(dtype);
+    *tensor_proto->mutable_tensor_shape() = tensor_shape;
+    tensor_proto->set_version_number(version_number);
+
+    return size_diff;
+  }
+};
+
+class ConstantSplitterFactory : public SizeSplitterFactory {
+ public:
+  using SizeSplitterFactory::SizeSplitterFactory;
+
+  absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
+      tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
+      std::vector<FieldType>* fields_in_parent, int size) override {
+    if (size < GetMaxSize()) return nullptr;
+    NodeDef* node = google::protobuf::DynamicCastMessage<NodeDef>(message);
+    if (node->op() != "Const")
+      return absl::UnimplementedError(absl::StrCat(
+          "Currently only able to split 'Const' nodes that are larger than the "
+          "2GB maximum proto size. Got node of type '",
+          node->op(), "' with size: ", size, "."));
+    ConstantSplitter* splitter =
+        new ConstantSplitter(message, parent_splitter, fields_in_parent);
+    return absl::WrapUnique(splitter);
+  }
+};
+
+class FunctionDefSplitter : public SizeSplitter {
+ public:
+  using SizeSplitter::SizeSplitter;
+  absl::StatusOr<int> BuildChunksReturnSize() override {
+    size_t current_size = GetInitialSize();
+    uint64_t max_size = GetMaxSize();
+    std::vector<FieldType> fields = {};
+    // First check if the entire FunctionDef can be split into a separate chunk.
+    // We do this before the `RepeatedMessageSplitter`, which is costly because
+    // it iterates through every `node_def`.
+    if (LARGE_SIZE_CHECK(current_size, max_size) && current_size < max_size) {
+      auto splitter = LargeNodeSplitter<FunctionDef>(message(), this, &fields);
+      splitter.SetInitialSize(current_size);
+      return splitter.BuildChunksReturnSize();
+    } else if (current_size > max_size) {
+      ConstantSplitterFactory constant_splitter_factory;
+      LargeNodeSplitterFactory<NodeDef> large_node_splitter_factory;
+      std::vector<SizeSplitterFactory*> factories = {
+          &constant_splitter_factory, &large_node_splitter_factory};
+      auto ret = RepeatedFieldSplitter<FunctionDef, NodeDef>::Create(
+          message(), this, &fields, "node_def"s, &factories);
+      if (!ret.ok()) return ret.status();
+      auto splitter = ret.value();
+      return splitter.BuildChunksReturnSize();
+    }
+
+    return 0;
+  }
+};
+
+class FunctionDefSplitterFactory : public SizeSplitterFactory {
+ public:
+  using SizeSplitterFactory::SizeSplitterFactory;
+
+  absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
+      tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
+      std::vector<FieldType>* fields_in_parent, int size) override {
+    FunctionDefSplitter* splitter =
+        new FunctionDefSplitter(message, parent_splitter, fields_in_parent);
+    return absl::WrapUnique(splitter);
+  }
+};
+
+}  // namespace
+
+absl::Status GraphDefSplitter::BuildChunks() {
+  TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
+  GraphDef* g = google::protobuf::DynamicCastMessage<GraphDef>(message());
+  uint64_t max_size = GetMaxSize();
+  size_t graph_size = GetInitialSize();
+
+  if (graph_size < max_size) return absl::OkStatus();
+
+  // Set up GraphDef.node and GraphDef.library.function splitters.
+  std::vector<FieldType> field_in_parent = {};
+  ConstantSplitterFactory constant_splitter_factory;
+  LargeNodeSplitterFactory<NodeDef> large_node_splitter_factory;
+  std::vector<SizeSplitterFactory*> factories = {&constant_splitter_factory,
+                                                 &large_node_splitter_factory};
+  auto node_splitter_ret = RepeatedFieldSplitter<GraphDef, NodeDef>::Create(
+      g, this, &field_in_parent, "node"s, &factories);
+  if (!node_splitter_ret.ok()) return node_splitter_ret.status();
+  auto node_splitter = node_splitter_ret.value();
+
+  FunctionDefSplitterFactory function_splitter_factory;
+  std::vector<FieldType> library_field = {"library"s};
+  std::vector<SizeSplitterFactory*> fn_factories = {&function_splitter_factory};
+  auto library_splitter_ret =
+      RepeatedFieldSplitter<FunctionDefLibrary, FunctionDef>::Create(
+          g->mutable_library(), this, &library_field, "function"s,
+          &fn_factories);
+  if (!library_splitter_ret.ok()) return library_splitter_ret.status();
+  auto library_splitter = library_splitter_ret.value();
+  size_t library_size = g->library().ByteSizeLong();
+  library_splitter.SetInitialSize(library_size);
+
+  size_t approx_node_size = graph_size - library_size;
+  node_splitter.SetInitialSize(approx_node_size);
+
+  // Call node and library splitters.
+  if (library_size > approx_node_size) {
+    TF_ASSIGN_OR_RETURN(int size_diff,
+                        library_splitter.BuildChunksReturnSize());
+    library_size -= size_diff;
+    if (approx_node_size + library_size > max_size) {
+      TF_ASSIGN_OR_RETURN(int size_diff, node_splitter.BuildChunksReturnSize());
+      approx_node_size -= size_diff;
+    }
+  } else {
+    TF_ASSIGN_OR_RETURN(int size_diff, node_splitter.BuildChunksReturnSize());
+    approx_node_size -= size_diff;
+    if (approx_node_size + library_size > max_size) {
+      TF_ASSIGN_OR_RETURN(int size_diff,
+                          library_splitter.BuildChunksReturnSize());
+      library_size -= size_diff;
+    }
+  }
+
+  // Recompute graph size. If the graph size is still greater than the max size,
+  // separate the entire library into a separate chunk.
+  if (g->ByteSizeLong() > max_size) {
+    LargeNodeSplitter<FunctionDefLibrary> entire_library_splitter(
+        g->mutable_library(), this, &library_field);
+    // The library chunk must be inserted before the function chunks generated
+    // by `library_splitter` which may have the same field tags.
+    int index = 1;
+    entire_library_splitter.SetChunkIndex(&index);
+    TF_RETURN_IF_ERROR(entire_library_splitter.BuildChunks());
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow::tools::proto_splitter

tensorflow/tools/proto_splitter/cc/graph_def_splitter.h

@@ -0,0 +1,35 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_GRAPH_DEF_SPLITTER_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_GRAPH_DEF_SPLITTER_H_
+
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// GraphDef Splitter.
+// Modifies the user-provided message when building chunks.
+class GraphDefSplitter : public ComposableSplitter {
+ public:
+  // Initializer.
+  using ComposableSplitter::ComposableSplitter;
+
+  absl::Status BuildChunks() override;
+};
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_GRAPH_DEF_SPLITTER_H_

tensorflow/tools/proto_splitter/cc/graph_def_splitter_test.cc

@@ -0,0 +1,299 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/graph_def_splitter.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <variant>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/cord.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/tools/proto_splitter/cc/max_size.h"
+#include "tensorflow/tools/proto_splitter/cc/test_util.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+namespace {
+
+using ::tensorflow::proto_splitter::ChunkedMessage;
+
+TEST(GraphDefSplitterTest, TestLargeConstant) {
+  GraphDef proto;
+  const std::string graph_def_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-large-constant.pb");
+  int64_t max_size = 500;
+  DebugSetMaxSize(max_size);
+
+  TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                           graph_def_path, &proto));
+  EXPECT_GE(proto.ByteSizeLong(), GetMaxSize());
+  std::string large_constant_1, large_constant_2;
+  const std::variant<std::string, absl::Cord>& tensor_constant_1 =
+      proto.node(2).attr().at("value").tensor().tensor_content();
+  const std::variant<std::string, absl::Cord>& tensor_constant_2 =
+      proto.node(4).attr().at("value").tensor().tensor_content();
+  if (std::holds_alternative<std::string>(tensor_constant_1)) {
+    large_constant_1 = std::get<std::string>(tensor_constant_1);
+  } else {
+    absl::CopyCordToString(std::get<absl::Cord>(tensor_constant_1),
+                           &large_constant_1);
+  }
+  if (std::holds_alternative<std::string>(tensor_constant_2)) {
+    large_constant_2 = std::get<std::string>(tensor_constant_2);
+  } else {
+    absl::CopyCordToString(std::get<absl::Cord>(tensor_constant_2),
+                           &large_constant_2);
+  }
+
+  GraphDefSplitter splitter(&proto);
+  TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
+  ChunkedMessage* chunked_message = x.chunked_message;
+  ASSERT_NE(chunked_message, nullptr);
+  EXPECT_THAT(*chunked_message,
+              EqualsProto(R"pb(chunk_index: 0
+                               chunked_fields {
+                                 field_tag { field: 1 }
+                                 field_tag { index: 2 }
+                                 field_tag { field: 5 }
+                                 field_tag { map_key { s: "value" } }
+                                 field_tag { field: 8 }
+                                 field_tag { field: 4 }
+                                 message { chunk_index: 1 }
+                               }
+                               chunked_fields {
+                                 field_tag { field: 1 }
+                                 field_tag { index: 4 }
+                                 field_tag { field: 5 }
+                                 field_tag { map_key { s: "value" } }
+                                 field_tag { field: 8 }
+                                 field_tag { field: 4 }
+                                 message { chunk_index: 2 }
+                               })pb"));
+
+  std::vector<MessageBytes>* chunks = x.chunks;
+  ASSERT_NE(chunks, nullptr);
+  EXPECT_CHUNK_SIZES(chunks, max_size);
+
+  EXPECT_THAT((*chunks)[1],
+              ::testing::VariantWith<std::string>(large_constant_1));
+  EXPECT_THAT((*chunks)[2],
+              ::testing::VariantWith<std::string>(large_constant_2));
+}
+
+TEST(GraphDefSplitterTest, TestLargeNodes) {
+  GraphDef proto;
+  const std::string graph_def_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-large-nodes.pb");
+  int64_t max_size = 200;
+  DebugSetMaxSize(max_size);
+
+  TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                           graph_def_path, &proto));
+  EXPECT_GE(proto.ByteSize(), GetMaxSize());
+
+  // Get the large nodes that should be extracted for later comparison.
+  NodeDef node_1 = proto.node(1);
+  NodeDef node_2 = proto.node(2);
+  NodeDef node_3 = proto.node(3);
+  NodeDef node_5 = proto.node(5);
+
+  GraphDefSplitter splitter(&proto);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
+  ChunkedMessage* chunked_message = x.chunked_message;
+  ASSERT_NE(chunked_message, nullptr);
+  EXPECT_THAT(*chunked_message, EqualsProto(R"pb(chunk_index: 0
+                                                 chunked_fields {
+                                                   field_tag { field: 1 }
+                                                   field_tag { index: 1 }
+                                                   message { chunk_index: 1 }
+                                                 }
+                                                 chunked_fields {
+                                                   field_tag { field: 1 }
+                                                   field_tag { index: 2 }
+                                                   message { chunk_index: 2 }
+                                                 }
+                                                 chunked_fields {
+                                                   field_tag { field: 1 }
+                                                   field_tag { index: 3 }
+                                                   message { chunk_index: 3 }
+                                                 }
+                                                 chunked_fields {
+                                                   field_tag { field: 1 }
+                                                   field_tag { index: 5 }
+                                                   message { chunk_index: 4 }
+                                                 })pb"));
+  std::vector<MessageBytes>* chunks = x.chunks;
+  ASSERT_NE(chunks, nullptr);
+  EXPECT_CHUNK_SIZES(chunks, max_size);
+
+  EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
+      (*chunks)[1]));
+  EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
+      (*chunks)[2]));
+  EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
+      (*chunks)[3]));
+  EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
+      (*chunks)[4]));
+
+  EXPECT_THAT(
+      *std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[1]).get(),
+      EqualsProto(node_1));
+  EXPECT_THAT(
+      *std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[2]).get(),
+      EqualsProto(node_2));
+  EXPECT_THAT(
+      *std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[3]).get(),
+      EqualsProto(node_3));
+  EXPECT_THAT(
+      *std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[4]).get(),
+      EqualsProto(node_5));
+}
+TEST(GraphDefSplitterTest, TestLotsNodes) {
+  GraphDef proto;
+  const std::string graph_def_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-lots-nodes.pb");
+
+  // split-lots-nodes.pb has 15 nodes that are 95 or 96 bytes each. The max size
+  // is set to "exactly" the size of 5 nodes, but with the extra encoding bytes,
+  // only 4 nodes should fit in each chunk. Thus, there should be exactly 4
+  // chunks created for all 15 nodes.
+  int64_t max_size = 96 * 5;
+  DebugSetMaxSize(max_size);
+
+  TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                           graph_def_path, &proto));
+  EXPECT_GE(proto.ByteSize(), GetMaxSize());
+  int expected_node_size = proto.node_size();
+
+  GraphDefSplitter splitter(&proto);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
+
+  ChunkedMessage* chunked_message = x.chunked_message;
+  ASSERT_NE(chunked_message, nullptr);
+  EXPECT_THAT(
+      *chunked_message,
+      EqualsProto(R"pb(chunk_index: 0
+                       chunked_fields { message { chunk_index: 1 } }
+                       chunked_fields { message { chunk_index: 2 } }
+                       chunked_fields { message { chunk_index: 3 } }
+                       chunked_fields { message { chunk_index: 4 } })pb"));
+
+  std::vector<MessageBytes>* chunks = x.chunks;
+  ASSERT_NE(chunks, nullptr);
+  EXPECT_CHUNK_SIZES(chunks, max_size);
+
+  int actual_node_size = 0;
+  for (MessageBytes& chunk : *chunks) {
+    GraphDef* message = nullptr;
+    if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
+            chunk)) {
+      message = google::protobuf::DynamicCastMessage<GraphDef>(
+          std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk).get());
+    } else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) {
+      message = google::protobuf::DynamicCastMessage<GraphDef>(
+          std::get<tsl::protobuf::Message*>(chunk));
+    } else {
+      EXPECT_FALSE(std::holds_alternative<std::string>(chunk));
+    }
+    actual_node_size += message->node_size();
+  }
+  EXPECT_EQ(actual_node_size, expected_node_size);
+}
+
+TEST(GraphDefSplitterTest, TestFunctionLotsOfNodes) {
+  GraphDef proto;
+  const std::string graph_def_path = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "tools/proto_splitter/testdata",
+      "function-lots-of-nodes.pb");
+  int64_t max_size = 500;
+  DebugSetMaxSize(max_size);
+
+  TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                           graph_def_path, &proto));
+  EXPECT_GE(proto.ByteSize(), GetMaxSize());
+
+  GraphDefSplitter splitter(&proto);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
+  std::vector<MessageBytes>* chunks = x.chunks;
+  ASSERT_NE(chunks, nullptr);
+  EXPECT_CHUNK_SIZES(chunks, max_size);
+}
+
+TEST(GraphDefSplitterTest, TestFunctionLargeNodes) {
+  GraphDef proto;
+  const std::string graph_def_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "function-large-nodes.pb");
+  int64_t max_size = 200;
+  DebugSetMaxSize(max_size);
+
+  TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                           graph_def_path, &proto));
+  EXPECT_GE(proto.ByteSize(), GetMaxSize());
+
+  GraphDefSplitter splitter(&proto);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
+  std::vector<MessageBytes>* chunks = x.chunks;
+  ASSERT_NE(chunks, nullptr);
+  EXPECT_CHUNK_SIZES(chunks, max_size);
+}
+
+TEST(GraphDefSplitterTest, TestGraphAndFunction) {
+  GraphDef proto;
+  const std::string graph_def_path = io::JoinPath(
+      testing::TensorFlowSrcRoot(), "tools/proto_splitter/testdata",
+      "graph-def-and-function.pb");
+  int64_t max_size = 200;
+  DebugSetMaxSize(max_size);
+
+  TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                           graph_def_path, &proto));
+  EXPECT_GE(proto.ByteSize(), GetMaxSize());
+
+  GraphDefSplitter splitter(&proto);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
+  std::vector<MessageBytes>* chunks = x.chunks;
+  ASSERT_NE(chunks, nullptr);
+  EXPECT_CHUNK_SIZES(chunks, max_size);
+
+  TF_ASSERT_OK(splitter.Write("/tmp/hoi"));
+}
+
+}  // namespace
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/cc/large_node_splitter.h

@@ -0,0 +1,84 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_LARGE_NODE_SPLITTER_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_LARGE_NODE_SPLITTER_H_
+
+#include <memory>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter_base.h"
+#include "tensorflow/tools/proto_splitter/cc/max_size.h"
+#include "tensorflow/tools/proto_splitter/cc/size_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tsl/platform/errors.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// Extracts messages that are large but not over the limit into a separate
+// message chunk.
+template <typename MessageType>
+class LargeNodeSplitter : public SizeSplitter {
+ public:
+  using SizeSplitter::SizeSplitter;
+
+  void SetChunkIndex(int* index) { index_ = index; }
+
+  absl::StatusOr<int> BuildChunksReturnSize() override;
+
+ private:
+  int* index_ = nullptr;
+};
+
+template <typename MessageType>
+absl::StatusOr<int> LargeNodeSplitter<MessageType>::BuildChunksReturnSize() {
+  MessageType* msg =
+      tsl::protobuf::DynamicCastToGenerated<MessageType>(message());
+  int initial_size = GetInitialSize();
+  std::shared_ptr<MessageType> new_msg = std::make_shared<MessageType>();
+  msg->Swap(new_msg.get());
+  std::vector<FieldType> fields = {};
+  auto x = std::make_unique<MessageBytes>(new_msg);
+  TF_RETURN_IF_ERROR(AddChunk(std::move(x), &fields, index_));
+  return initial_size;
+}
+
+template <typename MessageType>
+class LargeNodeSplitterFactory : public SizeSplitterFactory {
+ public:
+  using SizeSplitterFactory::SizeSplitterFactory;
+
+  absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
+      tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
+      std::vector<FieldType>* fields_in_parent, int size) override;
+};
+
+template <typename MessageType>
+absl::StatusOr<std::unique_ptr<SizeSplitter>>
+LargeNodeSplitterFactory<MessageType>::CreateSplitter(
+    tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
+    std::vector<FieldType>* fields_in_parent, int size) {
+  if (!(LARGE_SIZE_CHECK(size, GetMaxSize()))) return nullptr;
+  LargeNodeSplitter<MessageType>* splitter = new LargeNodeSplitter<MessageType>(
+      message, parent_splitter, fields_in_parent);
+  return absl::WrapUnique(splitter);
+}
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_LARGE_NODE_SPLITTER_H_

tensorflow/tools/proto_splitter/cc/max_size.cc

@@ -0,0 +1,32 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/max_size.h"
+
+#include "absl/synchronization/mutex.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+ABSL_CONST_INIT absl::Mutex global_mutex(absl::kConstInit);
+
+// The default max size is set to a bit less than 2GB, since the proto splitter
+// isn't extremely precise.
+uint64_t ProtoMaxSize = ((uint64_t)1 << 31) - 500;
+
+uint64_t GetMaxSize() { return ProtoMaxSize; }
+
+void DebugSetMaxSize(uint64_t size) { ProtoMaxSize = size; }
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/cc/max_size.h

@@ -0,0 +1,34 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_MAX_SIZE_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_MAX_SIZE_H_
+
+#include <cstdint>
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// Get max size allowed by each chunk in the proto splitter.
+uint64_t GetMaxSize();
+
+// Set the max size. Should only be used for testing purposes.
+void DebugSetMaxSize(uint64_t size);
+
+// Heuristic for determine whether to greedily generate a chunk.
+#define LARGE_SIZE_CHECK(size, max_size) size > max_size / 3
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_MAX_SIZE_H_

tensorflow/tools/proto_splitter/cc/repeated_field_splitter.h

@@ -0,0 +1,169 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_REPEATED_FIELD_SPLITTER_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_REPEATED_FIELD_SPLITTER_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/max_size.h"
+#include "tensorflow/tools/proto_splitter/cc/size_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow::tools::proto_splitter {
+
+// Splitter that works on repeated message fields.
+template <typename ParentMessage, typename RepeatedMessage>
+class RepeatedFieldSplitter : public SizeSplitter {
+ public:
+  static absl::StatusOr<RepeatedFieldSplitter> Create(
+      tsl::protobuf::Message* message, ComposableSplitter* parent_splitter,
+      std::vector<FieldType>* fields_in_parent, const FieldType& repeated_field,
+      std::vector<SizeSplitterFactory*>* splitter_factories);
+
+  absl::StatusOr<int> BuildChunksReturnSize() override;
+
+ private:
+  explicit RepeatedFieldSplitter(
+      tsl::protobuf::Message* message, ComposableSplitter* parent_splitter,
+      std::vector<FieldType>* fields_in_parent, const FieldType& repeated_field,
+      std::vector<SizeSplitterFactory*>* splitter_factories)
+      : SizeSplitter(message, parent_splitter, fields_in_parent),
+        repeated_field_(repeated_field),
+        splitter_factories_(splitter_factories) {}
+
+  FieldType repeated_field_;
+  std::vector<SizeSplitterFactory*>* splitter_factories_;
+};
+
+// Additional bytes added to each node to account for the extra info needed to
+// encode the field key (realistically 3 but making it 5 for some wiggle room).
+constexpr int kExtraBytes = 5;
+
+template <typename ParentMessage, typename RepeatedMessage>
+absl::StatusOr<RepeatedFieldSplitter<ParentMessage, RepeatedMessage>>
+RepeatedFieldSplitter<ParentMessage, RepeatedMessage>::Create(
+    tsl::protobuf::Message* message, ComposableSplitter* parent_splitter,
+    std::vector<FieldType>* fields_in_parent, const FieldType& repeated_field,
+    std::vector<SizeSplitterFactory*>* splitter_factories) {
+  TF_ASSIGN_OR_RETURN(auto field_ret, GetField(*message, {repeated_field}));
+  if (!field_ret.field->is_repeated()) {
+    return absl::FailedPreconditionError("Unable to split non-repeated field.");
+  }
+
+  auto ret = RepeatedFieldSplitter<ParentMessage, RepeatedMessage>(
+      message, parent_splitter, fields_in_parent, repeated_field,
+      splitter_factories);
+  return ret;
+}
+
+template <typename ParentMessage, typename RepeatedMessage>
+absl::StatusOr<int>
+RepeatedFieldSplitter<ParentMessage, RepeatedMessage>::BuildChunksReturnSize() {
+  TF_ASSIGN_OR_RETURN(MutableFieldResult mfr,
+                      GetMutableField(message(), {repeated_field_}));
+  tsl::protobuf::Message* parent = mfr.parent;
+  const tsl::protobuf::FieldDescriptor* repeated_field = mfr.field;
+
+  uint64_t max_size = GetMaxSize();
+  size_t initial_size = GetInitialSize();
+
+  // List of indices at which to split the repeated field. For example, [3, 5]
+  // means that the field list is split into: [:3], [3:5], [5:]
+  std::vector<int> repeated_msg_split;
+  // Track the total byte size of the current node split.
+  uint64_t total_size = 0;
+
+  // Linearly iterate through all nodes. It may be possible to optimize this
+  // further by making best guesses as to where to split the nodes, since
+  // most nodes (aside from constants) are relatively small.
+  int repeated_field_length =
+      parent->GetReflection()->FieldSize(*parent, repeated_field);
+  for (int i = 0; i < repeated_field_length; ++i) {
+    tsl::protobuf::Message* node =
+        parent->GetReflection()->MutableRepeatedMessage(parent, repeated_field,
+                                                        i);
+    auto node_size = node->ByteSizeLong();
+
+    std::vector<FieldType> new_fields = {repeated_field_, i};
+
+    for (auto factory : *splitter_factories_) {
+      TF_ASSIGN_OR_RETURN(
+          std::unique_ptr<SizeSplitter> new_splitter,
+          factory->CreateSplitter(node, this, &new_fields, node_size));
+      if (new_splitter != nullptr) {
+        TF_ASSIGN_OR_RETURN(auto size_diff,
+                            new_splitter->BuildChunksReturnSize());
+        node_size -= size_diff;
+      }
+    }
+    if (total_size + node_size > max_size) {
+      repeated_msg_split.push_back(i);
+      total_size = 0;
+    }
+    total_size += node_size + kExtraBytes;
+  }
+
+  if (!repeated_msg_split.empty()) {
+    auto repeated_nodes_ptrs =
+        parent->GetReflection()
+            ->template MutableRepeatedPtrField<RepeatedMessage>(parent,
+                                                                repeated_field);
+
+    std::vector<RepeatedMessage*> extracted_nodes(repeated_field_length);
+    repeated_nodes_ptrs->ExtractSubrange(0, repeated_field_length,
+                                         &extracted_nodes.at(0));
+    // Last range end is the size of the repeated field.
+    repeated_msg_split.push_back(repeated_field_length);
+
+    int range_start = 0;
+    for (int range_end : repeated_msg_split) {
+      auto new_msg = std::make_shared<ParentMessage>();
+      std::vector<FieldType> empty_fields;
+      auto x = std::make_unique<MessageBytes>(new_msg);
+      TF_RETURN_IF_ERROR(AddChunk(std::move(x), &empty_fields));
+
+      // Move nodes into new_msg.
+      TF_ASSIGN_OR_RETURN(auto new_ret,
+                          GetMutableField(new_msg.get(), repeated_field_));
+
+      for (int j = range_start; j < range_end; ++j) {
+        new_msg->GetReflection()->AddAllocatedMessage(
+            new_msg.get(), new_ret.field, extracted_nodes[j]);
+      }
+
+      range_start = range_end;
+    }
+  }
+
+  // Estimate the size diff by subtracting the first computed chunk size from
+  // the initial size of the repeated field.
+  return initial_size - message()->ByteSizeLong();
+}
+
+}  // namespace tensorflow::tools::proto_splitter
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_REPEATED_FIELD_SPLITTER_H_

tensorflow/tools/proto_splitter/cc/saved_model_splitter.cc

@@ -0,0 +1,62 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/saved_model_splitter.h"
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/tools/proto_splitter/cc/graph_def_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/large_node_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/max_size.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// Required in OSS to prevent string to bool conversion in FieldType variant.
+using namespace std::string_literals;  // NOLINT
+
+absl::Status SavedModelSplitter::BuildChunks() {
+  TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
+  SavedModel* sm = google::protobuf::DynamicCastMessage<SavedModel>(message());
+  int max_size = GetMaxSize();
+  if (GetInitialSize() < max_size) return absl::OkStatus();
+
+  std::vector<FieldType> fields_to_graph_def = {"meta_graphs"s, 0,
+                                                "graph_def"s};
+  GraphDefSplitter graph_def_splitter(
+      sm->mutable_meta_graphs(0)->mutable_graph_def(), this,
+      &fields_to_graph_def);
+  TF_RETURN_IF_ERROR(graph_def_splitter.BuildChunks());
+
+  if (sm->ByteSizeLong() < max_size) return absl::OkStatus();
+
+  LargeNodeSplitter<GraphDef> entire_graph_splitter(
+      sm->mutable_meta_graphs(0)->mutable_graph_def(), this,
+      &fields_to_graph_def);
+  int index = 1;
+  entire_graph_splitter.SetChunkIndex(&index);
+  TF_RETURN_IF_ERROR(entire_graph_splitter.BuildChunks());
+
+  return absl::OkStatus();
+}
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/cc/saved_model_splitter.h

@@ -0,0 +1,35 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_SAVED_MODEL_SPLITTER_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SAVED_MODEL_SPLITTER_H_
+
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// SavedModel Splitter.
+// Modifies the user-provided message when building chunks.
+class SavedModelSplitter : public ComposableSplitter {
+ public:
+  // Initializer.
+  using ComposableSplitter::ComposableSplitter;
+
+  absl::Status BuildChunks() override;
+};
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SAVED_MODEL_SPLITTER_H_

tensorflow/tools/proto_splitter/cc/saved_model_splitter_test.cc

@@ -0,0 +1,89 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/saved_model_splitter.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <variant>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/tools/proto_splitter/cc/max_size.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+namespace {
+
+// Ensures that all Messages are less than the max size. std::string chunks are
+// not limited by the max size, so they are ignored in this check.
+#define EXPECT_CHUNK_SIZES(chunks, max_size)                                \
+  do {                                                                      \
+    for (auto chunk : *chunks) {                                            \
+      if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(  \
+              chunk)) {                                                     \
+        EXPECT_LE(std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk)  \
+                      ->ByteSizeLong(),                                     \
+                  max_size);                                                \
+      } else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) {  \
+        EXPECT_LE(std::get<tsl::protobuf::Message*>(chunk)->ByteSizeLong(), \
+                  max_size);                                                \
+      }                                                                     \
+    }                                                                       \
+  } while (0)
+
+std::string NonChunkedSavedModel() {
+  return io::JoinPath(testing::TensorFlowSrcRoot(), "cc", "saved_model",
+                      "testdata", "chunked_saved_model", "non_chunked_model",
+                      "saved_model.pb");
+}
+
+TEST(SavedModelSplitterTest, TestSplit) {
+  SavedModel proto;
+  int64_t max_size = 80000;
+  DebugSetMaxSize(max_size);
+
+  TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
+                                           NonChunkedSavedModel(), &proto));
+  EXPECT_GE(proto.ByteSizeLong(), GetMaxSize());
+
+  SavedModelSplitter splitter(&proto);
+
+  TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
+  std::vector<MessageBytes>* chunks = x.chunks;
+  ASSERT_NE(chunks, nullptr);
+
+  // Should create a new chunk with the single large constant.
+  EXPECT_EQ(2, chunks->size());
+  EXPECT_CHUNK_SIZES(chunks, max_size);
+}
+
+}  // namespace
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/cc/size_splitter.h

@@ -0,0 +1,59 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_SIZE_SPLITTER_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SIZE_SPLITTER_H_
+
+#include <memory>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// Abstract composable Splitter that returns the size of chunks created.
+// Modifies the user-provided message when building chunks.
+class SizeSplitter : public ComposableSplitter {
+ public:
+  // Initializer.
+  using ComposableSplitter::ComposableSplitter;
+
+  absl::Status BuildChunks() override {
+    return BuildChunksReturnSize().status();
+  }
+
+  // Chunks the message and returns the int size diff.
+  virtual absl::StatusOr<int> BuildChunksReturnSize() = 0;
+};
+
+class SizeSplitterFactory {
+ public:
+  explicit SizeSplitterFactory() = default;
+
+  // Creates a new SizeSplitter object based on the input parameters. May
+  // return nullptr if no split is necessary.
+  virtual absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
+      tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
+      std::vector<FieldType>* fields_in_parent, int size) = 0;
+
+  virtual ~SizeSplitterFactory() = default;
+};
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SIZE_SPLITTER_H_

tensorflow/tools/proto_splitter/cc/split.h

@@ -0,0 +1,60 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_SPLIT_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SPLIT_H_
+
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tensorflow/tools/proto_splitter/versions.pb.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+using ::tensorflow::proto_splitter::ChunkedMessage;
+using ::tensorflow::proto_splitter::VersionDef;
+
+// Interface for proto message splitters.
+class Splitter {
+ public:
+  virtual ~Splitter() = default;
+
+  // Split message into chunks.
+  virtual absl::StatusOr<ChunkedProto> Split() = 0;
+
+  // Write message to disk and returns the full output path.
+  virtual absl::StatusOr<std::string> Write(std::string file_prefix) = 0;
+
+  // Version info about the Splitter and required Merger versions.
+  virtual VersionDef Version() = 0;
+
+ protected:
+  Splitter() = default;
+  Splitter(const Splitter&) = default;
+  Splitter(Splitter&&) = default;
+};
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SPLIT_H_

tensorflow/tools/proto_splitter/cc/test_util.h

@@ -0,0 +1,136 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_TEST_UTIL_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_TEST_UTIL_H_
+
+#include <ostream>
+#include <string>
+
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+// Ensures that all Messages are less than the max size. std::string chunks are
+// not limited by the max size, so they are ignored in this check.
+#define EXPECT_CHUNK_SIZES(chunks, max_size)                                \
+  do {                                                                      \
+    for (auto chunk : *chunks) {                                            \
+      if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(  \
+              chunk)) {                                                     \
+        EXPECT_LE(std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk)  \
+                      ->ByteSizeLong(),                                     \
+                  max_size);                                                \
+      } else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) {  \
+        EXPECT_LE(std::get<tsl::protobuf::Message*>(chunk)->ByteSizeLong(), \
+                  max_size);                                                \
+      }                                                                     \
+    }                                                                       \
+  } while (0)
+
+inline std::string SerializeAsString(const tsl::protobuf::Message& message) {
+  std::string result;
+  {
+    // Use a nested block to guarantee triggering coded_stream's destructor
+    // before `result` is used. Due to copy elision, this code works without
+    // the nested block, but small, innocent looking changes can break it.
+    tsl::protobuf::io::StringOutputStream string_stream(&result);
+    tsl::protobuf::io::CodedOutputStream coded_stream(&string_stream);
+    coded_stream.SetSerializationDeterministic(true);
+    message.SerializeToCodedStream(&coded_stream);
+  }
+  return result;
+}
+
+template <typename MessageType>
+tsl::StatusOr<MessageType> ParseTextProto(const char* text_proto) {
+  tsl::protobuf::TextFormat::Parser parser;
+  MessageType parsed_proto;
+  bool success =
+      tsl::protobuf::TextFormat::ParseFromString(text_proto, &parsed_proto);
+  if (!success) {
+    return absl::InvalidArgumentError(
+        "Input text could not be parsed to proto.");
+  }
+  return parsed_proto;
+}
+
+// TODO(b/229726259): EqualsProto is not available in OSS.
+
+// Simple implementation of a proto matcher comparing string representations.
+// This will fail if the string representations are not deterministic.
+class ProtoStringMatcher {
+ public:
+  explicit ProtoStringMatcher(const tsl::protobuf::Message& expected)
+      : expected_(SerializeAsString(expected)) {}
+
+  explicit ProtoStringMatcher(const std::string& expected)
+      : text_format_(expected), use_text_format_(true) {}
+
+  template <typename Message>
+  bool MatchAndExplain(const Message& p,
+                       ::testing::MatchResultListener*) const {
+    if (use_text_format_) {
+      Message* m = p.New();
+      CHECK(tsl::protobuf::TextFormat::ParseFromString(text_format_, m))
+          << "Failed to parse proto text: " << text_format_;
+      std::string expected = SerializeAsString(*m);
+      delete m;
+      return SerializeAsString(p) == expected;
+    } else {
+      return SerializeAsString(p) == expected_;
+    }
+  }
+
+  void DescribeTo(::std::ostream* os) const {
+    if (use_text_format_)
+      *os << text_format_;
+    else
+      *os << expected_;
+  }
+  void DescribeNegationTo(::std::ostream* os) const {
+    *os << "not equal to expected message: ";
+    auto out_message = use_text_format_ ? &text_format_ : &expected_;
+    if (out_message->size() < 1e5)
+      *os << *out_message;
+    else
+      *os << "(too large to print) \n";
+  }
+
+ private:
+  const std::string expected_;
+  const std::string text_format_;
+  const bool use_text_format_ = false;
+};
+
+inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
+    const tsl::protobuf::Message& x) {
+  return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
+}
+
+inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
+    const std::string& x) {
+  return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
+}
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_TEST_UTIL_H_

tensorflow/tools/proto_splitter/cc/util.cc

@@ -0,0 +1,805 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/util.h"
+
+#include <cstdint>
+#include <functional>
+#include <ios>
+#include <optional>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "riegeli/base/maker.h"  // from @riegeli
+#include "riegeli/base/types.h"  // from @riegeli
+#include "riegeli/bytes/fd_reader.h"  // from @riegeli
+#include "riegeli/bytes/string_reader.h"  // from @riegeli
+#include "riegeli/records/record_reader.h"  // from @riegeli
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system_helper.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+using ::tensorflow::proto_splitter::ChunkedField;
+
+namespace {
+absl::StatusOr<int> FieldInt(const FieldType& field) {
+  switch (field.index()) {
+    case 0: {  // std::string
+      int value;
+      if (absl::SimpleAtoi(std::get<std::string>(field), &value)) return value;
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Unable to convert '", std::get<std::string>(field), "' to int."));
+      break;
+    }
+    case 1:  // int
+      return std::get<int>(field);
+    case 2:  // bool
+      return std::get<bool>(field) ? 1 : 0;
+    default:
+      // This should not happen.
+      return absl::InvalidArgumentError("Invalid field type.");
+  }
+}
+absl::StatusOr<bool> FieldBool(const FieldType& field) {
+  switch (field.index()) {
+    case 0: {  // std::string
+      std::string s = std::get<std::string>(field);
+      if (s == "true") return true;
+      if (s == "false") return false;
+      return absl::InvalidArgumentError("Unable to convert '" + s +
+                                        "' to bool.");
+    }
+    case 1:  // int
+      return std::get<int>(field) != 0;
+    case 2:  // bool
+      return std::get<bool>(field);
+    default:
+      // This should not happen.
+      return absl::InvalidArgumentError("Invalid field type.");
+  }
+}
+std::string FieldString(const FieldType& field) {
+  switch (field.index()) {
+    case 0:  // std::string
+      return std::get<std::string>(field);
+    case 1:  // int
+      return std::to_string(std::get<int>(field));
+    case 2:  // bool
+      return std::get<bool>(field) ? "true" : "false";
+    default:
+      // This should not happen.
+      CHECK(false) << "Invalid field type: " << field.index();  // Crash OK.
+      return "";
+  }
+}
+
+template <typename FieldCallback, typename MapKeyCallback,
+          typename ListIndexCallback>
+absl::Status WalkFields(const tsl::protobuf::Descriptor& desc,
+                        const std::vector<FieldType>& fields,
+                        FieldCallback field_callback,
+                        MapKeyCallback map_key_callback,
+                        ListIndexCallback list_index_callback) {
+  const tsl::protobuf::Descriptor* message_desc = &desc;
+  std::vector<FieldType>::const_iterator it = fields.begin();
+
+  while (it != fields.end()) {
+    const tsl::protobuf::FieldDescriptor* field_desc;
+    if (std::holds_alternative<std::string>(*it)) {
+      field_desc = message_desc->FindFieldByName(std::get<std::string>(*it));
+    } else if (std::holds_alternative<int>(*it)) {
+      field_desc = message_desc->FindFieldByNumber(std::get<int>(*it));
+    } else {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Invalid field, expected int or str: ", FieldString(*it)));
+    }
+
+    if (field_desc == nullptr) {
+      return absl::NotFoundError(
+          absl::StrCat("Field not found: ", FieldString(*it)));
+    }
+    TF_RETURN_IF_ERROR(field_callback(field_desc));
+    message_desc = field_desc->message_type();
+
+    it++;
+    // Handle special field types (map key and list index).
+    if (it != fields.end()) {
+      if (field_desc->is_map()) {
+        auto key_field = message_desc->FindFieldByNumber(1);
+        auto status = map_key_callback(*key_field, *it);
+        if (!status.ok()) {
+          return status;
+        }
+        // The next element in `fields` must be a field in the value message.
+        message_desc = message_desc->FindFieldByName("value")->message_type();
+        it++;
+      } else if (field_desc->is_repeated()) {
+        auto status = list_index_callback(*it);
+        if (!status.ok()) {
+          return status;
+        }
+        it++;
+      }
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status AddMapKey(const tsl::protobuf::FieldDescriptor& key_field,
+                       const FieldType& map_key, ChunkedField& chunked_field) {
+  auto msg = chunked_field.add_field_tag()->mutable_map_key();
+  switch (key_field.type()) {
+    case tsl::protobuf::FieldDescriptor::TYPE_BOOL: {
+      TF_ASSIGN_OR_RETURN(auto key, FieldBool(map_key));
+      msg->set_boolean(key);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::TYPE_STRING: {
+      msg->set_s(FieldString(map_key));
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::TYPE_INT32: {
+      TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+      msg->set_i32(key);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::TYPE_INT64: {
+      TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+      msg->set_i64(key);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::TYPE_UINT32: {
+      TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+      msg->set_ui32(key);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::TYPE_UINT64: {
+      TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+      msg->set_ui64(key);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::TYPE_FIXED64:
+    case tsl::protobuf::FieldDescriptor::TYPE_FIXED32:
+    case tsl::protobuf::FieldDescriptor::TYPE_SINT32:
+    case tsl::protobuf::FieldDescriptor::TYPE_SINT64:
+    case tsl::protobuf::FieldDescriptor::TYPE_SFIXED32:
+    case tsl::protobuf::FieldDescriptor::TYPE_SFIXED64:
+      return absl::UnimplementedError(
+          "Fixed map key types not implemented yet.");
+    case tsl::protobuf::FieldDescriptor::TYPE_DOUBLE:
+    case tsl::protobuf::FieldDescriptor::TYPE_FLOAT:
+    case tsl::protobuf::FieldDescriptor::TYPE_ENUM:
+    case tsl::protobuf::FieldDescriptor::TYPE_BYTES:
+    case tsl::protobuf::FieldDescriptor::TYPE_MESSAGE:
+    case tsl::protobuf::FieldDescriptor::TYPE_GROUP:
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Encountered type that is not supported for MapKey.type: ",
+          key_field.type()));
+
+    default:
+      return absl::FailedPreconditionError("Proto field type not recognized.");
+  }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<FieldType> GetMapKeyFromFieldIndex(
+    ::tensorflow::proto_splitter::FieldIndex field_index) {
+  if (!field_index.has_map_key())
+    return absl::FailedPreconditionError(
+        "Field index doesn't contain a map key.");
+
+  switch (field_index.map_key().type_case()) {
+    case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kBoolean:
+      return field_index.map_key().boolean();
+      break;
+    case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kS:
+      return field_index.map_key().s();
+      break;
+    case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kI32:
+      return field_index.map_key().i32();
+      break;
+    case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kI64:
+      // Cast to int type, which may be lossy. We'll deal with it when it
+      // becomes an issue.
+      return static_cast<int>(field_index.map_key().i64());
+      break;
+    case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kUi32:
+      return static_cast<int>(field_index.map_key().ui32());
+      break;
+    case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kUi64:
+      return static_cast<int>(field_index.map_key().ui64());
+      break;
+    case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
+        TYPE_NOT_SET:
+    default:
+      return absl::FailedPreconditionError(
+          absl::StrCat("Unknown map key type: ", field_index.DebugString()));
+  }
+}
+
+}  // namespace
+
+absl::StatusOr<const std::vector<Field>> GetFieldTypes(
+    const tsl::protobuf::RepeatedPtrField<
+        ::tensorflow::proto_splitter::FieldIndex>& field_tags) {
+  std::vector<Field> fields;
+  for (int fti = 0; fti < field_tags.size();) {
+    switch (field_tags[fti].kind_case()) {
+      case ::tensorflow::proto_splitter::FieldIndex::KindCase::kField:
+        fields.push_back(
+            Field(static_cast<int>(field_tags[fti].field()), std::nullopt));
+        fti++;
+        if (fti == field_tags.size()) break;
+        // Multiple field tags may correspond to a single field when the field
+        // is repeated or a map.
+        if (field_tags[fti].has_index()) {
+          fields.back().second = static_cast<int>(field_tags[fti++].index());
+        } else if (field_tags[fti].has_map_key()) {
+          TF_ASSIGN_OR_RETURN(const FieldType& map_key,
+                              GetMapKeyFromFieldIndex(field_tags[fti++]));
+          fields.back().second = map_key;
+        }
+        break;
+      case ::tensorflow::proto_splitter::FieldIndex::KindCase::kIndex:
+        return absl::FailedPreconditionError(
+            "Index doesn't belong to any field.");
+        break;
+      case ::tensorflow::proto_splitter::FieldIndex::KindCase::kMapKey:
+        return absl::FailedPreconditionError(
+            "Map key doesn't belong to any field.");
+        break;
+      case ::tensorflow::proto_splitter::FieldIndex::KindCase::KIND_NOT_SET:
+      default:
+        return absl::FailedPreconditionError(absl::StrCat(
+            "Unknown field kind: ", field_tags[fti].DebugString()));
+    }
+  }
+  return fields;
+}
+
+absl::Status SetRepeatedFieldElement(
+    tsl::protobuf::Message* message,
+    const tsl::protobuf::FieldDescriptor* field_desc, uint64_t field_index,
+    const std::string& chunk,
+    std::function<absl::Status(void)> message_callback) {
+  if (field_desc->is_map())
+    return absl::FailedPreconditionError("Field is a map.");
+  const tsl::protobuf::Reflection* reflection = message->GetReflection();
+  if (field_index >= reflection->FieldSize(*message, field_desc))
+    return absl::OutOfRangeError(
+        absl::StrCat("Field index out of range: ", field_index));
+
+  switch (field_desc->cpp_type()) {
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32:
+      reflection->SetRepeatedInt32(message, field_desc, std::stoi(chunk),
+                                   field_index);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64:
+      reflection->SetRepeatedInt64(message, field_desc, std::stoll(chunk),
+                                   field_index);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32:
+      reflection->SetRepeatedUInt32(message, field_desc, std::stoul(chunk),
+                                    field_index);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64:
+      reflection->SetRepeatedUInt64(message, field_desc, std::stoull(chunk),
+                                    field_index);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_DOUBLE:
+      reflection->SetRepeatedDouble(message, field_desc, std::stod(chunk),
+                                    field_index);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_FLOAT:
+      reflection->SetRepeatedFloat(message, field_desc, std::stof(chunk),
+                                   field_index);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL:
+      bool b;
+      std::istringstream(chunk) >> std::boolalpha >> b;
+      reflection->SetRepeatedBool(message, field_desc, b, field_index);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_ENUM:
+      reflection->SetRepeatedEnum(
+          message, field_desc, field_index,
+          field_desc->enum_type()->FindValueByName(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING:
+      reflection->SetRepeatedString(message, field_desc, field_index, chunk);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_MESSAGE:
+      return message_callback();
+      break;
+    default:
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Proto field type not recognized: ", field_desc->cpp_type_name()));
+      break;
+  }
+  return absl::OkStatus();
+}
+
+absl::Status SetFieldElement(
+    tsl::protobuf::Message* message,
+    const tsl::protobuf::FieldDescriptor* field_desc, const std::string& chunk,
+    std::function<absl::Status(void)> message_callback) {
+  const tsl::protobuf::Reflection* reflection = message->GetReflection();
+
+  switch (field_desc->cpp_type()) {
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32:
+      reflection->SetInt32(message, field_desc, std::stoi(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64:
+      reflection->SetInt64(message, field_desc, std::stoll(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32:
+      reflection->SetUInt32(message, field_desc, std::stoul(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64:
+      reflection->SetUInt64(message, field_desc, std::stoull(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_DOUBLE:
+      reflection->SetDouble(message, field_desc, std::stod(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_FLOAT:
+      reflection->SetFloat(message, field_desc, std::stof(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL:
+      bool b;
+      std::istringstream(chunk) >> std::boolalpha >> b;
+      reflection->SetBool(message, field_desc, b);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_ENUM:
+      reflection->SetEnum(message, field_desc,
+                          field_desc->enum_type()->FindValueByName(chunk));
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING:
+      reflection->SetString(message, field_desc, chunk);
+      break;
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_MESSAGE:
+      return message_callback();
+      break;
+    default:
+      return absl::FailedPreconditionError("Proto field type not recognized.");
+      break;
+  }
+  return absl::OkStatus();
+}
+
+absl::Status AddMapEntry(tsl::protobuf::Message* message,
+                         const tsl::protobuf::FieldDescriptor* field_desc,
+                         FieldType map_key) {
+  tsl::protobuf::Message* map_entry =
+      message->GetReflection()->AddMessage(message, field_desc);
+  const tsl::protobuf::Reflection* reflection = map_entry->GetReflection();
+  const tsl::protobuf::FieldDescriptor* key =
+      field_desc->message_type()->FindFieldByNumber(1);
+  switch (key->cpp_type()) {
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32: {
+      TF_ASSIGN_OR_RETURN(auto map_key_int32, FieldInt(map_key));
+      reflection->SetInt32(map_entry, key, map_key_int32);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64: {
+      TF_ASSIGN_OR_RETURN(auto map_key_int64, FieldInt(map_key));
+      reflection->SetInt64(map_entry, key, map_key_int64);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32: {
+      TF_ASSIGN_OR_RETURN(auto map_key_uint32, FieldInt(map_key));
+      reflection->SetUInt32(map_entry, key, map_key_uint32);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64: {
+      TF_ASSIGN_OR_RETURN(auto map_key_uint64, FieldInt(map_key));
+      reflection->SetUInt64(map_entry, key, map_key_uint64);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL: {
+      TF_ASSIGN_OR_RETURN(auto map_key_bool, FieldBool(map_key));
+      reflection->SetBool(map_entry, key, map_key_bool);
+      break;
+    }
+    case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING: {
+      reflection->SetString(map_entry, key, FieldString(map_key));
+      break;
+    }
+    default:
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Proto field type not recognized.", key->cpp_type_name()));
+  }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<MutableFieldResult> GetMutableField(
+    tsl::protobuf::Message* message, const std::vector<FieldType>& fields) {
+  tsl::protobuf::Message* parent = message;
+  const tsl::protobuf::FieldDescriptor* field = nullptr;
+  int index = -1;
+
+  auto field_callback =
+      [&parent, &field, &index](
+          const tsl::protobuf::FieldDescriptor* field_desc) -> absl::Status {
+    const tsl::protobuf::Reflection* reflection = parent->GetReflection();
+    if (field != nullptr) {
+      // Update the parent proto using the previous FieldDescriptor.
+      if (field->is_map()) {
+        parent = reflection->MutableRepeatedMessage(parent, field, index);
+        parent = parent->GetReflection()->MutableMessage(
+            parent, parent->GetDescriptor()->FindFieldByNumber(2));
+        index = -1;
+      } else if (field->is_repeated()) {
+        parent = reflection->MutableRepeatedMessage(parent, field, index);
+        index = -1;
+      } else {
+        parent = reflection->MutableMessage(parent, field);
+      }
+    }
+
+    field = field_desc;
+    return absl::OkStatus();
+  };
+
+  auto map_key_callback = [&parent, &field, &index](
+                              const tsl::protobuf::FieldDescriptor& key_field,
+                              const FieldType& map_key) -> absl::Status {
+    // Finds the index of the map field w/ the corresponding map_key.
+    const tsl::protobuf::Reflection* reflection = parent->GetReflection();
+
+    TF_ASSIGN_OR_RETURN(int found_index,
+                        FindMapKey(*parent, *field, &key_field, map_key));
+
+    if (found_index == -1) {
+      return absl::NotFoundError(
+          absl::StrCat("Error when getting map field, couldn't find key: ",
+                       FieldString(map_key)));
+    }
+    if (reflection->FieldSize(*parent, field) <= found_index) {
+      return absl::NotFoundError(absl::StrCat(
+          "Can't access index ", index, " in field '", field->name(),
+          "' (size = ", reflection->FieldSize(*parent, field), ")."));
+    }
+    index = found_index;
+    return absl::OkStatus();
+  };
+  auto list_index_callback = [&parent, &field,
+                              &index](FieldType list_index) -> absl::Status {
+    TF_ASSIGN_OR_RETURN(index, FieldInt(list_index));
+    const tsl::protobuf::Reflection* reflection = parent->GetReflection();
+    if (reflection->FieldSize(*parent, field) <= index) {
+      return absl::NotFoundError(absl::StrCat(
+          "Can't access index ", index, " in field '", field->name(),
+          "' (size = ", reflection->FieldSize(*parent, field), ")."));
+    }
+    return absl::OkStatus();
+  };
+  TF_RETURN_IF_ERROR(WalkFields(*message->GetDescriptor(), fields,
+                                field_callback, map_key_callback,
+                                list_index_callback));
+
+  MutableFieldResult result{.parent = parent, .field = field, .index = index};
+  return result;
+}
+
+absl::StatusOr<MutableFieldResult> GetMutableField(
+    tsl::protobuf::Message* message, const Field& field) {
+  std::vector<FieldType> fields = {field.first};
+  if (field.second != std::nullopt) fields.push_back(field.second.value());
+  return GetMutableField(message, fields);
+}
+
+absl::StatusOr<MutableFieldResult> GetMutableField(
+    tsl::protobuf::Message* message, const FieldType& field_type) {
+  std::vector<FieldType> fields = {field_type};
+  return GetMutableField(message, fields);
+}
+
+absl::StatusOr<FieldResult> GetField(const tsl::protobuf::Message& message,
+                                     const std::vector<FieldType>& fields) {
+  const tsl::protobuf::Message* parent = &message;
+  const tsl::protobuf::FieldDescriptor* field = nullptr;
+  int index = -1;
+
+  auto field_callback =
+      [&parent, &field, &index](
+          const tsl::protobuf::FieldDescriptor* field_desc) -> absl::Status {
+    const tsl::protobuf::Reflection* reflection = parent->GetReflection();
+    if (field != nullptr) {
+      // Update the parent proto using the previous FieldDescriptor.
+      if (field->is_map()) {
+        parent = &reflection->GetRepeatedMessage(*parent, field, index);
+        parent = &parent->GetReflection()->GetMessage(
+            *parent, parent->GetDescriptor()->FindFieldByNumber(2));
+        index = -1;
+      } else if (field->is_repeated()) {
+        parent = &reflection->GetRepeatedMessage(*parent, field, index);
+        index = -1;
+      } else {
+        parent = &reflection->GetMessage(*parent, field);
+      }
+    }
+
+    field = field_desc;
+    return absl::OkStatus();
+  };
+
+  auto map_key_callback = [&parent, &field, &index](
+                              const tsl::protobuf::FieldDescriptor& key_field,
+                              const FieldType& map_key) -> absl::Status {
+    // Finds the index of the map field w/ the corresponding map_key.
+    const tsl::protobuf::Reflection* reflection = parent->GetReflection();
+
+    TF_ASSIGN_OR_RETURN(int found_index,
+                        FindMapKey(*parent, *field, &key_field, map_key));
+
+    if (found_index == -1) {
+      return absl::NotFoundError(
+          absl::StrCat("Error when getting map field, couldn't find key: ",
+                       FieldString(map_key)));
+    }
+    if (reflection->FieldSize(*parent, field) <= found_index) {
+      return absl::NotFoundError(absl::StrCat(
+          "Can't access index ", index, " in field '", field->name(),
+          "' (size = ", reflection->FieldSize(*parent, field), ")."));
+    }
+    index = found_index;
+    return absl::OkStatus();
+  };
+  auto list_index_callback = [&parent, &field,
+                              &index](FieldType list_index) -> absl::Status {
+    TF_ASSIGN_OR_RETURN(index, FieldInt(list_index));
+    const tsl::protobuf::Reflection* reflection = parent->GetReflection();
+    if (reflection->FieldSize(*parent, field) <= index) {
+      return absl::NotFoundError(absl::StrCat(
+          "Can't access index ", index, " in field '", field->name(),
+          "' (size = ", reflection->FieldSize(*parent, field), ")."));
+    }
+    return absl::OkStatus();
+  };
+  TF_RETURN_IF_ERROR(WalkFields(*message.GetDescriptor(), fields,
+                                field_callback, map_key_callback,
+                                list_index_callback));
+
+  FieldResult result{
+      .parent = parent,
+      .field = field,
+      .index = index,
+  };
+  return result;
+}
+
+absl::Status AddFieldTag(const tsl::protobuf::Descriptor& desc,
+                         const std::vector<FieldType>& fields,
+                         ChunkedField& chunked_field) {
+  auto field_callback =
+      [&chunked_field](
+          const tsl::protobuf::FieldDescriptor* field_desc) -> absl::Status {
+    chunked_field.add_field_tag()->set_field(field_desc->number());
+    return absl::OkStatus();
+  };
+  auto map_key_callback = [&chunked_field](
+                              const tsl::protobuf::FieldDescriptor& key_field,
+                              FieldType map_key) -> absl::Status {
+    return AddMapKey(key_field, map_key, chunked_field);
+  };
+  auto list_index_callback =
+      [&chunked_field](FieldType list_index) -> absl::Status {
+    TF_ASSIGN_OR_RETURN(auto index, FieldInt(list_index));
+    chunked_field.add_field_tag()->set_index(index);
+    return absl::OkStatus();
+  };
+
+  return WalkFields(desc, fields, field_callback, map_key_callback,
+                    list_index_callback);
+}
+
+absl::Status AddFieldTag(const tsl::protobuf::Descriptor& desc,
+                         const Field& field, ChunkedField& chunked_field) {
+  std::vector<FieldType> fields = {field.first};
+  if (field.second != std::nullopt) fields.push_back(field.second.value());
+  return AddFieldTag(desc, fields, chunked_field);
+}
+
+absl::StatusOr<int> FindMapKey(const tsl::protobuf::Message& parent,
+                               const tsl::protobuf::FieldDescriptor& map_field,
+                               const tsl::protobuf::FieldDescriptor* key_field,
+                               FieldType map_key) {
+  const tsl::protobuf::Reflection* reflection = parent.GetReflection();
+
+  if (!map_field.is_map()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "FindMapKey() was given a non map field: ", map_field.full_name()));
+  }
+
+  if (key_field == nullptr) {
+    key_field = map_field.message_type()->FindFieldByNumber(1);
+  }
+
+  bool found = false;
+  for (int i = 0; i < reflection->FieldSize(parent, &map_field) && !found;
+       i++) {
+    auto kv_pair = &reflection->GetRepeatedMessage(parent, &map_field, i);
+    auto kv_reflection = kv_pair->GetReflection();
+    switch (key_field->cpp_type()) {
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32: {
+        TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+        if (kv_reflection->GetInt32(*kv_pair, key_field) == key) {
+          found = true;
+        }
+        break;
+      }
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64: {
+        TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+        if (kv_reflection->GetInt64(*kv_pair, key_field) == key) {
+          found = true;
+        }
+        break;
+      }
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32: {
+        TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+        if (kv_reflection->GetUInt32(*kv_pair, key_field) == key) {
+          found = true;
+        }
+        break;
+      }
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64: {
+        TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
+        if (kv_reflection->GetUInt64(*kv_pair, key_field) == key) {
+          found = true;
+        }
+        break;
+      }
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL: {
+        TF_ASSIGN_OR_RETURN(auto key, FieldBool(map_key));
+        if (kv_reflection->GetBool(*kv_pair, key_field) == key) {
+          found = true;
+        }
+        break;
+      }
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING: {
+        if (kv_reflection->GetString(*kv_pair, key_field) ==
+            FieldString(map_key)) {
+          found = true;
+        }
+        break;
+      }
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_DOUBLE:
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_FLOAT:
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_ENUM:
+      case tsl::protobuf::FieldDescriptor::CPPTYPE_MESSAGE:
+        return absl::FailedPreconditionError(absl::StrCat(
+            "Type is not supported for MapKey.type: ", key_field->cpp_type()));
+      default:
+        return absl::FailedPreconditionError(
+            "Proto field type not recognized.");
+    }
+    if (found) return i;
+  }
+  return -1;
+}
+
+namespace {
+
+const std::vector<std::pair<double, std::string>>& ByteUnits() {
+  static const auto& byte_units =
+      *(new std::vector<std::pair<double, std::string>>(
+          {{1, "B"}, {1 << 10, "KiB"}, {1 << 20, "MiB"}, {1 << 30, "GiB"}}));
+  return byte_units;
+}
+
+}  // namespace
+
+std::string HumanReadableBytes(int64_t byte_count) {
+  for (int i = 1; i < ByteUnits().size(); i++) {
+    if (byte_count < ByteUnits()[i].first) {
+      return absl::StrFormat("%.1f%s", byte_count / ByteUnits()[i - 1].first,
+                             ByteUnits()[i - 1].second);
+    }
+  }
+  int i = ByteUnits().size() - 1;
+  return absl::StrFormat("%.1f%s", byte_count / ByteUnits()[i].first,
+                         ByteUnits()[i].second);
+}
+
+std::string HumanReadableDuration(int64_t microseconds) {
+  if (microseconds < 1000) {
+    return absl::StrFormat("%d microseconds", microseconds);
+  } else if (microseconds < 1e6) {
+    return absl::StrFormat("%.2f ms", microseconds / 1e3);
+  } else {
+    return absl::StrFormat("%.2f s", microseconds / 1e6);
+  }
+}
+
+absl::StatusOr<riegeli::RecordReader<riegeli::FdReader<>>> GetRiegeliReader(
+    absl::string_view cpb_file) {
+  riegeli::RecordReader reader(riegeli::Maker<riegeli::FdReader>(cpb_file));
+  if (!reader.ok()) {
+    return reader.status();
+  }
+  return reader;
+}
+
+absl::StatusOr<riegeli::RecordReader<riegeli::StringReader<>>>
+GetRiegeliStringReader(absl::string_view cpb_data) {
+  riegeli::RecordReader reader(
+      riegeli::Maker<riegeli::StringReader<>>(cpb_data));
+  if (!reader.ok()) {
+    return reader.status();
+  }
+  return reader;
+}
+
+absl::StatusOr<::tensorflow::proto_splitter::ChunkMetadata> GetChunkMetadata(
+    riegeli::RecordReaderBase& reader) {
+  ::tensorflow::proto_splitter::ChunkMetadata chunk_metadata;
+  bool read_metadata_success = reader.Seek(reader.Size().value()) &&
+                               reader.SeekBack() &&
+                               reader.ReadRecord(chunk_metadata);
+  if (read_metadata_success) return chunk_metadata;
+  return reader.status();
+}
+
+absl::StatusOr<std::string> ReadChunk(
+    riegeli::RecordReaderBase& reader,
+    const ::tensorflow::proto_splitter::ChunkInfo& chunk_info) {
+  riegeli::Position pos = chunk_info.offset();
+  std::string chunk(chunk_info.size(), '\0');
+  if (reader.Seek(pos) && reader.ReadRecord(chunk)) return chunk;
+  return absl::NotFoundError(
+      absl::StrCat("Chunk could not be found at position: ", pos, ".\n",
+                   reader.status().ToString()));
+}
+
+absl::StatusOr<bool> OnlyContainsPb(absl::string_view prefix) {
+  const std::string pb_file = absl::StrCat(prefix, ".pb");
+  const std::string cpb_file = absl::StrCat(prefix, ".cpb");
+  TF_ASSIGN_OR_RETURN(bool is_pb,
+                      internal::FileExists(Env::Default(), pb_file));
+  TF_ASSIGN_OR_RETURN(bool is_cpb,
+                      internal::FileExists(Env::Default(), cpb_file));
+  if (is_pb && !is_cpb) {
+    return true;
+  } else if (!is_pb && !is_cpb) {
+    return absl::NotFoundError(
+        absl::StrCat("Could not find SavedModel .pb or .cpb at supplied "
+                     "export directory path with prefix: ",
+                     prefix,
+                     ". Check that "
+                     "the directory exists and that you have the right "
+                     "permissions for accessing it."));
+  }
+  LOG(INFO) << "Reading chunked proto from " << cpb_file;
+  return false;
+}
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/cc/util.h

@@ -0,0 +1,171 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_CC_UTIL_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "riegeli/bytes/fd_reader.h"  // from @riegeli
+#include "riegeli/bytes/string_reader.h"  // from @riegeli
+#include "riegeli/records/record_reader.h"  // from @riegeli
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+
+using MessageBytes = std::variant<std::shared_ptr<tsl::protobuf::Message>,
+                                  tsl::protobuf::Message*, std::string>;
+
+struct ChunkedProto {
+  std::vector<MessageBytes>* chunks = nullptr;
+  ::tensorflow::proto_splitter::ChunkedMessage* chunked_message = nullptr;
+};
+
+// TODO(b/282796592): Consider switching to `tsl::protobuf::FieldPath` in the
+// future.
+
+// Fields can be represented using their name (string) or number (int). Map key
+// and repeated field index types are also included in this variant type.
+using FieldType = std::variant<std::string, int, bool>;
+using Field = std::pair<FieldType, std::optional<FieldType>>;
+
+// Convert a sequence of field tags to a vector of fields. A single field is a
+// std::vector<FieldType>, since multiple field tags may correspond to a single
+// field when the field is repeated or a map.
+absl::StatusOr<const std::vector<Field>> GetFieldTypes(
+    const tsl::protobuf::RepeatedPtrField<
+        ::tensorflow::proto_splitter::FieldIndex>& field_tags);
+
+// Sets message.field_desc[field_index] to the data contained in chunk,
+// according to the (cpp) type described by field_desc. Uses message_callback
+// (instead of simply assigning) when field_desc describes a message.
+absl::Status SetRepeatedFieldElement(
+    tsl::protobuf::Message* message,
+    const tsl::protobuf::FieldDescriptor* field_desc, uint64_t field_index,
+    const std::string& chunk,
+    std::function<absl::Status(void)> message_callback);
+
+// Sets message.field_desc to the data contained in chunk, according to the
+// (cpp) type described by field_desc. Uses message_callback (instead of simply
+// assigning) when field_desc describes a message.
+absl::Status SetFieldElement(
+    tsl::protobuf::Message* message,
+    const tsl::protobuf::FieldDescriptor* field_desc, const std::string& chunk,
+    std::function<absl::Status(void)> message_callback);
+
+// Adds a new map entry (repeated message element with key/value fields) to
+// message.field_desc (a map). The new map entry's key is set using map_key,
+// according to its type.
+absl::Status AddMapEntry(tsl::protobuf::Message* message,
+                         const tsl::protobuf::FieldDescriptor* field_desc,
+                         FieldType map_key);
+
+// Struct returned by `GetMutableField`. The field can be retrieved by using the
+// Reflection API on the parent.
+struct MutableFieldResult {
+  tsl::protobuf::Message* parent;
+  const tsl::protobuf::FieldDescriptor* field;
+  // If field is repeated or map, `index` is set to the list index or the
+  // position at which the map key appears. If the field is not repeated,
+  // `index` is -1.
+  int index;
+};
+
+// Returns a mutable parent, field descriptor, and int index in the case of a
+// repeated or map value field (or -1 if a non-repeated/map field).
+absl::StatusOr<MutableFieldResult> GetMutableField(
+    tsl::protobuf::Message* message, const std::vector<FieldType>& fields);
+
+absl::StatusOr<MutableFieldResult> GetMutableField(
+    tsl::protobuf::Message* message, const Field& field);
+
+// Gets info about the mutable field that's directly attached to message.
+absl::StatusOr<MutableFieldResult> GetMutableField(
+    tsl::protobuf::Message* message, const FieldType& field_type);
+
+// Struct returned by `GetField`. The field can be retrieved by using the
+// Reflection API on the parent.
+struct FieldResult {
+  const tsl::protobuf::Message* parent;
+  const tsl::protobuf::FieldDescriptor* field;
+  // If field is repeated or map, `index` is set to the list index or the
+  // position at which the map key appears. If the field is not repeated,
+  // `index` is -1.
+  int index;
+};
+
+// Returns the parent message, field descriptor, and int index from following
+// the provided message and fields.
+absl::StatusOr<FieldResult> GetField(const tsl::protobuf::Message& message,
+                                     const std::vector<FieldType>& fields);
+
+// Updates `field_tag` in the ChunkedField proto.
+absl::Status AddFieldTag(
+    const tsl::protobuf::Descriptor& desc, const std::vector<FieldType>& fields,
+    ::tensorflow::proto_splitter::ChunkedField& chunked_field);
+
+absl::Status AddFieldTag(
+    const tsl::protobuf::Descriptor& desc, const Field& field,
+    ::tensorflow::proto_splitter::ChunkedField& chunked_field);
+
+// Returns the index of the map key in the map field. If the key is not found,
+// returns -1.
+absl::StatusOr<int> FindMapKey(const tsl::protobuf::Message& parent,
+                               const tsl::protobuf::FieldDescriptor& map_field,
+                               const tsl::protobuf::FieldDescriptor* key_field,
+                               FieldType map_key);
+
+// Formats bytes into something more readable. (e.g. 52428800 -> "50.0MiB")
+std::string HumanReadableBytes(int64_t byte_count);
+// Formats microseconds into a more readable string.
+std::string HumanReadableDuration(int64_t microseconds);
+
+// Construct a reader object to read in records from the .cpb file.
+absl::StatusOr<riegeli::RecordReader<riegeli::FdReader<>>> GetRiegeliReader(
+    absl::string_view cpb_file);
+
+// Construct a reader object to read in records from a string.
+absl::StatusOr<riegeli::RecordReader<riegeli::StringReader<>>>
+GetRiegeliStringReader(absl::string_view cpb_data);
+
+// Read the last chunk, which contains metadata necessary for reading the
+// remaining chunks.
+absl::StatusOr<::tensorflow::proto_splitter::ChunkMetadata> GetChunkMetadata(
+    riegeli::RecordReaderBase& reader);
+
+// Use the `reader` to read in the chunk specified by `chunk_info`.
+absl::StatusOr<std::string> ReadChunk(
+    riegeli::RecordReaderBase& reader,
+    const ::tensorflow::proto_splitter::ChunkInfo& chunk_info);
+
+// Returns true if prefix can only be found as a .pb file, and false if a .cpb
+// file exists. Returns an error if neither .pb nor .cpb exist.
+absl::StatusOr<bool> OnlyContainsPb(absl::string_view prefix);
+
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_

tensorflow/tools/proto_splitter/cc/util_test.cc

@@ -0,0 +1,550 @@
+/* Copyright 2023 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/tools/proto_splitter/cc/util.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/tools/proto_splitter/cc/test_util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/status.h"
+#include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace tools::proto_splitter {
+namespace {
+
+using ::tensorflow::proto_splitter::ChunkedField;
+using ::tensorflow::proto_splitter_testdata::ManyFields;
+using ::testing::HasSubstr;
+using tsl::testing::IsOkAndHolds;
+using tsl::testing::StatusIs;
+
+// Required in OSS to prevent string to bool conversion in FieldType variant.
+using namespace std::string_literals;  // NOLINT
+
+absl::StatusOr<ManyFields> MakeManyFields() {
+  return ParseTextProto<ManyFields>(
+      R"pb(field_one {
+             repeated_field {}
+             repeated_field {
+               string_field: "inner_inner_string"
+               map_field_uint32 { key: 324 value: "map_value_324" }
+               map_field_uint32 { key: 543 value: "map_value_543" }
+             }
+           }
+           map_field_int64 { key: -1345 value: "map_value_-1345" }
+           nested_map_bool {
+             key: false
+             value { string_field: "string_false" }
+           }
+           nested_map_bool {
+             key: true
+             value { string_field: "string_true" }
+           })pb");
+}
+
+absl::StatusOr<
+    tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
+MakeFieldTags() {
+  TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
+                        field_tag { field: 2 }
+                        field_tag { index: 1505 }
+                        field_tag { field: 5 }
+                        field_tag { map_key { ui32: 123 } }
+                      )pb"));
+  return ret.field_tag();
+}
+
+absl::StatusOr<
+    tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
+MakeFieldTagsTooManyIndices() {
+  TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
+                        field_tag { field: 2 }
+                        field_tag { index: 1505 }
+                        field_tag { index: 1506 }
+                        field_tag { field: 5 }
+                        field_tag { map_key { ui32: 123 } }
+                      )pb"));
+  return ret.field_tag();
+}
+
+absl::StatusOr<
+    tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
+MakeFieldTagsTooManyMapKeys() {
+  TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
+                        field_tag { field: 2 }
+                        field_tag { index: 1505 }
+                        field_tag { field: 5 }
+                        field_tag { map_key { ui32: 123 } }
+                        field_tag { map_key: { ui32: 124 } }
+                      )pb"));
+  return ret.field_tag();
+}
+
+absl::StatusOr<
+    tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
+MakeFieldTagsMisplacedIndex() {
+  TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
+                        field_tag { field: 2 }
+                        field_tag { index: 1505 }
+                        field_tag { field: 5 }
+                        field_tag { map_key { ui32: 123 } }
+                        field_tag { index: 1504 }
+                      )pb"));
+  return ret.field_tag();
+}
+
+absl::StatusOr<
+    tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
+MakeFieldTagsMisplacedMapKey() {
+  TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
+                        field_tag { field: 2 }
+                        field_tag { index: 1505 }
+                        field_tag { map_key: { ui32: 321 } }
+                        field_tag { field: 5 }
+                        field_tag { map_key { ui32: 123 } }
+                      )pb"));
+  return ret.field_tag();
+}
+
+TEST(UtilTest, TestFieldTag) {
+  ManyFields message;
+  ChunkedField field;
+  std::vector<FieldType> fields = {"nested_map_bool"s, true, 2, 50, 7, false};
+  TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
+
+  EXPECT_THAT(field,
+              EqualsProto(R"pb(field_tag { field: 7 }
+                               field_tag { map_key { boolean: true } }
+                               field_tag { field: 2 }
+                               field_tag { index: 50 }
+                               field_tag { field: 7 }
+                               field_tag { map_key { boolean: false } })pb"));
+}
+TEST(UtilTest, TestFieldTagWithInt64MapKey) {
+  ManyFields message;
+  ChunkedField field;
+  std::vector<FieldType> fields = {"map_field_int64"s, -4234};
+  TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
+
+  EXPECT_THAT(field, EqualsProto(R"pb(field_tag { field: 6 }
+                                      field_tag { map_key { i64: -4234 } }
+              )pb"));
+}
+
+TEST(UtilTest, TestFieldTagWithUInt32MapKey) {
+  ManyFields message;
+  ChunkedField field;
+  std::vector<FieldType> fields = {"repeated_field"s, 1505, "map_field_uint32"s,
+                                   123};
+  TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
+
+  EXPECT_THAT(field, EqualsProto(R"pb(field_tag { field: 2 }
+                                      field_tag { index: 1505 }
+                                      field_tag { field: 5 }
+                                      field_tag { map_key { ui32: 123 } }
+              )pb"));
+}
+
+TEST(UtilTest, TestFieldTagConversion) {
+  ManyFields message;
+  ChunkedField field;
+  std::vector<FieldType> fields = {"repeated_field"s, "1505"s,
+                                   "map_field_uint32"s, 123};
+  TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
+
+  EXPECT_THAT(field, EqualsProto(R"pb(field_tag { field: 2 }
+                                      field_tag { index: 1505 }
+                                      field_tag { field: 5 }
+                                      field_tag { map_key { ui32: 123 } }
+              )pb"));
+}
+
+TEST(UtilTest, TestGetFieldTypes) {
+  TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTags());
+  EXPECT_THAT(
+      GetFieldTypes(tags),
+      absl_testing::IsOkAndHolds(std::vector<Field>{{2, 1505}, {5, 123}}));
+}
+
+TEST(UtilTest, TestGetFieldTypesThenAddFieldTags) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  ChunkedField chunked_field;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTags());
+  TF_ASSERT_OK_AND_ASSIGN(std::vector<Field> fields, GetFieldTypes(tags));
+  for (const auto& field : fields) {
+    TF_ASSERT_OK(AddFieldTag(*message.descriptor(), field, chunked_field));
+  }
+
+  EXPECT_THAT(chunked_field,
+              EqualsProto(R"pb(field_tag { field: 2 }
+                               field_tag { index: 1505 }
+                               field_tag { field: 5 }
+                               field_tag { map_key { ui32: 123 } }
+              )pb"));
+}
+
+TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsTooManyIndices) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  ChunkedField chunked_field;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsTooManyIndices());
+  EXPECT_THAT(
+      GetFieldTypes(tags),
+      absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
+                             HasSubstr("Index doesn't belong to any field")));
+}
+
+TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsTooManyMapKeys) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  ChunkedField chunked_field;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsTooManyMapKeys());
+  EXPECT_THAT(
+      GetFieldTypes(tags),
+      absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
+                             HasSubstr("Map key doesn't belong to any field")));
+}
+
+TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsMisplacedIndex) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  ChunkedField chunked_field;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsMisplacedIndex());
+  EXPECT_THAT(
+      GetFieldTypes(tags),
+      absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
+                             HasSubstr("Index doesn't belong to any field")));
+}
+
+TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsMisplacedMapKey) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  ChunkedField chunked_field;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsMisplacedMapKey());
+  EXPECT_THAT(
+      GetFieldTypes(tags),
+      absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
+                             HasSubstr("Map key doesn't belong to any field")));
+}
+
+TEST(UtilTest, TestSetRepeatedFieldElement) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message,
+      ParseTextProto<ManyFields>(R"pb(repeated_string_field: ""
+                                      repeated_string_field: "")pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("repeated_string_field");
+  const std::vector<std::string> chunks = {"repeated_string_one",
+                                           "repeated_string_two"};
+  TF_ASSERT_OK(SetRepeatedFieldElement(
+      &message, field_desc, 0, chunks[0],
+      []() -> absl::Status { return absl::OkStatus(); }));
+  TF_ASSERT_OK(SetRepeatedFieldElement(
+      &message, field_desc, 1, chunks[1],
+      []() -> absl::Status { return absl::OkStatus(); }));
+
+  EXPECT_FALSE(message.repeated_string_field().empty());
+  EXPECT_EQ(message.repeated_string_field().at(0), "repeated_string_one");
+  EXPECT_EQ(message.repeated_string_field().at(1), "repeated_string_two");
+}
+
+TEST(UtilTest, TestSetRepeatedFieldElementInvalidIndex) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message,
+      ParseTextProto<ManyFields>(R"pb(repeated_string_field: "")pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("repeated_string_field");
+
+  EXPECT_THAT(SetRepeatedFieldElement(
+                  &message, field_desc, 1, "",
+                  []() -> absl::Status { return absl::OkStatus(); }),
+              absl_testing::StatusIs(absl::StatusCode::kOutOfRange,
+                                     HasSubstr("Field index out of range")));
+}
+
+TEST(UtilTest, TestSetRepeatedFieldElementAlreadyExists) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message, ParseTextProto<ManyFields>(
+                        R"pb(repeated_string_field: "existing_string")pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("repeated_string_field");
+  const std::string chunk = "inner_inner_string_v2";
+
+  TF_EXPECT_OK(SetRepeatedFieldElement(
+      &message, field_desc, 0, chunk,
+      []() -> absl::Status { return absl::OkStatus(); }));
+
+  EXPECT_FALSE(message.repeated_string_field().empty());
+  EXPECT_NE(message.repeated_string_field().at(0), "existing_string");
+  EXPECT_EQ(message.repeated_string_field().at(0), "inner_inner_string_v2");
+}
+
+TEST(UtilTest, TestSetRepeatedFieldElementBadFieldDesc) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message, ParseTextProto<ManyFields>(
+                        R"pb(repeated_string_field: "existing_string")pb"));
+
+  const tsl::protobuf::FieldDescriptor* singular_field_desc =
+      message.GetDescriptor()->FindFieldByName("string_field");
+  EXPECT_DEATH(
+      auto status = SetRepeatedFieldElement(
+          &message, singular_field_desc, 0, "",
+          []() -> absl::Status { return absl::OkStatus(); }),
+      HasSubstr("Field is singular; the method requires a repeated field"));
+
+  const tsl::protobuf::FieldDescriptor* map_field_desc =
+      message.GetDescriptor()->FindFieldByName("map_field_uint32");
+  EXPECT_THAT(SetRepeatedFieldElement(
+                  &message, map_field_desc, 0, "",
+                  []() -> absl::Status { return absl::OkStatus(); }),
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
+                                     HasSubstr("Field is a map")));
+}
+
+TEST(UtilTest, TestSetRepeatedFieldElementMessage) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message, ParseTextProto<ManyFields>(R"pb(repeated_field {})pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("repeated_field");
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto chunk_proto, ParseTextProto<ManyFields>(R"pb(string_field: "")pb"));
+
+  const std::string chunk = chunk_proto.SerializeAsString();
+
+  TF_ASSERT_OK(SetRepeatedFieldElement(
+      &message, field_desc, 0, chunk,
+      [&message, &field_desc]() -> absl::Status {
+        tsl::protobuf::Message* inner_message =
+            message.GetReflection()->MutableRepeatedMessage(&message,
+                                                            field_desc, 0);
+        inner_message->GetReflection()->SetString(
+            inner_message,
+            field_desc->message_type()->FindFieldByName("string_field"),
+            "inner_string");
+        return absl::OkStatus();
+      }));
+
+  EXPECT_FALSE(message.repeated_field().empty());
+  EXPECT_EQ(message.repeated_field().at(0).string_field(), "inner_string");
+}
+
+TEST(UtilTest, TestSetFieldElement) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message, ParseTextProto<ManyFields>(R"pb(string_field: "")pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("string_field");
+  const std::string chunk = "string_field_v2";
+  TF_ASSERT_OK(
+      SetFieldElement(&message, field_desc, chunk,
+                      []() -> absl::Status { return absl::OkStatus(); }));
+
+  EXPECT_EQ(message.string_field(), "string_field_v2");
+}
+
+TEST(UtilTest, TestSetFieldElementInvalidRepeated) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message,
+      ParseTextProto<ManyFields>(R"pb(repeated_string_field: "")pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("repeated_string_field");
+
+  EXPECT_DEATH(
+      auto status =
+          SetFieldElement(&message, field_desc, "",
+                          []() -> absl::Status { return absl::OkStatus(); }),
+      HasSubstr("Field is repeated; the method requires a singular field"));
+}
+
+TEST(UtilTest, TestSetFieldElementAlreadyExists) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto message,
+      ParseTextProto<ManyFields>(R"pb(string_field: "existing_string")pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("string_field");
+  const std::string chunk = "inner_inner_string_v2";
+
+  TF_EXPECT_OK(
+      SetFieldElement(&message, field_desc, chunk,
+                      []() -> absl::Status { return absl::OkStatus(); }));
+
+  EXPECT_EQ(message.string_field(), "inner_inner_string_v2");
+}
+
+TEST(UtilTest, TestSetFieldElementMessage) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message,
+                          ParseTextProto<ManyFields>(R"pb(field_one {})pb"));
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("field_one");
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto chunk_proto, ParseTextProto<ManyFields>(R"pb(string_field: "")pb"));
+
+  const std::string chunk = chunk_proto.SerializeAsString();
+
+  TF_ASSERT_OK(SetFieldElement(
+      &message, field_desc, chunk, [&message, &field_desc]() -> absl::Status {
+        tsl::protobuf::Message* inner_message =
+            message.GetReflection()->MutableMessage(&message, field_desc);
+        inner_message->GetReflection()->SetString(
+            inner_message,
+            field_desc->message_type()->FindFieldByName("string_field"),
+            "inner_string");
+        return absl::OkStatus();
+      }));
+
+  EXPECT_TRUE(message.has_field_one());
+  EXPECT_EQ(message.field_one().string_field(), "inner_string");
+}
+
+TEST(UtilTest, TestAddMapEntry) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  const tsl::protobuf::FieldDescriptor* field_desc =
+      message.GetDescriptor()->FindFieldByName("map_field_int64");
+  FieldType map_key = -4234;
+  TF_ASSERT_OK(AddMapEntry(&message, field_desc, map_key));
+
+  TF_ASSERT_OK_AND_ASSIGN(int map_entry_index,
+                          FindMapKey(message, *field_desc, nullptr, map_key));
+  tsl::protobuf::Message* map_entry =
+      message.GetReflection()->MutableRepeatedMessage(&message, field_desc,
+                                                      map_entry_index);
+  map_entry->GetReflection()->SetString(
+      map_entry, field_desc->message_type()->FindFieldByNumber(2),
+      "map_value_-4234");
+
+  EXPECT_EQ(message.map_field_int64().at(-4234), "map_value_-4234");
+}
+
+TEST(UtilTest, GetFieldInvalidIndex) {
+  std::vector<FieldType> fields = {"field_one"s, "repeated_field"s, 100};
+
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  EXPECT_THAT(GetField(message, fields),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound,
+                                     HasSubstr("Can't access index 100")));
+}
+
+TEST(UtilTest, GetFieldInvalidField) {
+  std::vector<FieldType> fields = {"field_one"s, "INVALID"s};
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  EXPECT_THAT(GetField(message, fields),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound,
+                                     HasSubstr("Field not found: INVALID")));
+}
+
+TEST(UtilTest, GetFieldInvalidMapKey) {
+  std::vector<FieldType> fields = {"map_field_int64"s, 10000};
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  EXPECT_THAT(GetField(message, fields),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound,
+                                     HasSubstr("couldn't find key: 10000")));
+}
+TEST(UtilTest, GetField) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  std::vector<FieldType> fields = {"field_one"s, "repeated_field"s};
+
+  TF_ASSERT_OK_AND_ASSIGN(auto ret, GetField(message, fields));
+  EXPECT_EQ(-1, ret.index);
+  EXPECT_EQ(2, ret.parent->GetReflection()->FieldSize(*ret.parent, ret.field));
+
+  fields = {"nested_map_bool"s, "true"s, "string_field"s};
+  TF_ASSERT_OK_AND_ASSIGN(auto ret2, GetField(message, fields));
+  EXPECT_EQ(-1, ret2.index);
+  EXPECT_EQ("string_true",
+            ret2.parent->GetReflection()->GetString(*ret2.parent, ret2.field));
+}
+
+TEST(UtilTest, FindMapKey) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  const tsl::protobuf::FieldDescriptor* map_field =
+      message.GetDescriptor()->FindFieldByName("nested_map_bool");
+  TF_ASSERT_OK_AND_ASSIGN(int i,
+                          FindMapKey(message, *map_field, nullptr, true));
+  EXPECT_EQ(1, i);
+  TF_ASSERT_OK_AND_ASSIGN(i, FindMapKey(message, *map_field, nullptr, false));
+  EXPECT_EQ(0, i);
+
+  map_field = message.GetDescriptor()->FindFieldByName("map_field_int64");
+  TF_ASSERT_OK_AND_ASSIGN(i, FindMapKey(message, *map_field, nullptr, -1345));
+  EXPECT_EQ(0, i);
+  TF_ASSERT_OK_AND_ASSIGN(i, FindMapKey(message, *map_field, nullptr, 1345));
+  EXPECT_EQ(-1, i);
+}
+
+TEST(UtilTest, FindMapKeyInvalid) {
+  TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
+  const tsl::protobuf::FieldDescriptor* not_map_field =
+      message.GetDescriptor()->FindFieldByName("string_field");
+  EXPECT_THAT(FindMapKey(message, *not_map_field, nullptr, -1345),
+              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                     HasSubstr("was given a non map field")));
+}
+
+TEST(UtilTest, TestHumanReadableBytes) {
+  EXPECT_EQ("1.0KiB", HumanReadableBytes(1024));
+  EXPECT_EQ("5.5KiB", HumanReadableBytes(5632));
+  EXPECT_EQ("52.2KiB", HumanReadableBytes(53432));
+  EXPECT_EQ("72.9MiB", HumanReadableBytes(76493281));
+  EXPECT_EQ("57.0MiB", HumanReadableBytes(5.977e7));
+  EXPECT_EQ("1.0GiB", HumanReadableBytes(1.074e9));
+  EXPECT_EQ("15.4GiB", HumanReadableBytes(16493342281));
+}
+
+TEST(UtilTest, TestHumanReadableDuration) {
+  EXPECT_EQ("534 microseconds", HumanReadableDuration(534));
+  EXPECT_EQ("1.00 ms", HumanReadableDuration(1000));
+  EXPECT_EQ("14.33 ms", HumanReadableDuration(14328));
+  EXPECT_EQ("95.83 s", HumanReadableDuration(95825433));
+}
+
+TEST(UtilTest, TestReadChunk) {
+  std::string cpb_file =
+      io::JoinPath(tensorflow::testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-standard.cpb");
+
+  TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
+
+  auto read_metadata = GetChunkMetadata(reader);
+  if (!read_metadata.ok()) {
+    reader.Close();
+    TF_ASSERT_OK(read_metadata.status());
+  }
+  ::tensorflow::proto_splitter::ChunkMetadata metadata = read_metadata.value();
+  std::vector<::tensorflow::proto_splitter::ChunkInfo> chunks_info(
+      metadata.chunks().begin(), metadata.chunks().end());
+
+  for (const auto& chunk_info : chunks_info) {
+    TF_ASSERT_OK_AND_ASSIGN(std::string chunk, ReadChunk(reader, chunk_info));
+    ASSERT_EQ(chunk.size(), chunk_info.size());
+  }
+}
+
+}  // namespace
+}  // namespace tools::proto_splitter
+}  // namespace tensorflow

tensorflow/tools/proto_splitter/chunk.proto

@@ -0,0 +1,92 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.proto_splitter;
+
+import "tensorflow/tools/proto_splitter/versions.proto";
+
+option cc_enable_arenas = true;
+
+// The ChunkMetadata is generated and saved as the last chunk when exported with
+// Splitter.write(prefix). This metadata is always located in the {prefix}.cpb
+// file.
+message ChunkMetadata {
+  // The version of the Splitter was used (for information purposes only), and
+  // what versions of the Merger can be used to read the chunked proto.
+  VersionDef version = 1;
+  // Metadata about how/where to read the chunk data.
+  repeated ChunkInfo chunks = 2;
+  // Recursive structural information of the proto that was passed to the
+  // Splitter. Used to merge the chunked data into a single proto.
+  ChunkedMessage message = 3;
+}
+
+message ChunkInfo {
+  // List of types supported by the Merger.
+  enum Type {
+    UNSET = 0;
+    MESSAGE = 1;
+    BYTES = 2;
+  }
+  Type type = 1;
+  uint64 size = 2;    // Size of this chunk in bytes.
+  uint64 offset = 3;  // Byte offset of this chunk.
+}
+
+message ChunkedMessage {
+  // Index of the ChunkInfo in ChunkMetadata.chunks.
+  optional uint64 chunk_index = 1;
+
+  // The same field can appear multiple times in this list. When this
+  // happens, the parsed chunk data is merged in the order that they appear.
+  // The way the field is merged depends on the type of field. e.g. Repeated
+  // fields are merged by appending them, bytes are merged concatenated.
+  repeated ChunkedField chunked_fields = 2;
+}
+
+message ChunkedField {
+  // Stores the nested field numbers, which tell us where and how to merge
+  // this chunk.
+  // E.g. If the parent ChunkedMessage contains a GraphDef, this field value
+  // could point to: `GraphDef.node[15].attr["value"].tensor.float_value`
+  // Field tag = [] (empty list) indicates that the chunked message is the same
+  // type as the parent. This is useful for slicing repeated fields.
+  repeated FieldIndex field_tag = 1;
+  ChunkedMessage message = 3;  // This allows us to build a tree of chunked
+                               // messages.
+}
+
+message FieldIndex {
+  message MapKey {
+    oneof type {
+      string s = 1;
+      bool boolean = 2;
+      uint32 ui32 = 3;
+      uint64 ui64 = 4;
+      int32 i32 = 5;
+      int64 i64 = 6;
+      // TODO(278326592): Add sfixed key types if needed.
+    }
+  }
+  oneof kind {
+    // uint32 is the maximum represent field tags:
+    // https://stackoverflow.com/a/63666587
+    uint32 field = 1;  // Field number
+    MapKey map_key = 2;
+    uint64 index = 3;
+  }
+}

tensorflow/tools/proto_splitter/constants.py

@@ -0,0 +1,37 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Common constants used by proto splitter modules."""
+
+# The splitter algorithm isn't extremely precise, so the max is set to a little
+# less than 2GB.
+#
+# TODO: b/380463192 - Consider fixing the split algorithm to handle edge cases
+# accurately and raising the max size to 2GB.
+_MAX_SIZE = 1 << 30
+
+
+def debug_set_max_size(value: int) -> None:
+  """Sets the max size allowed for each proto chunk (used for debugging only).
+
+  Args:
+    value: int byte size
+  """
+  global _MAX_SIZE
+  _MAX_SIZE = value
+
+
+def max_size() -> int:
+  """Returns the maximum size each proto chunk."""
+  return _MAX_SIZE

tensorflow/tools/proto_splitter/g3doc/in-depth-guide.md

@@ -0,0 +1,492 @@
+# Proto Splitter / Merger Library
+
+This doc lists implementation details about the [Proto Splitter/Merger library](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/tools/proto_splitter). New Splitters should take these details into consideration to generate valid chunks and metadata that are compatible with the Merger. If you'd just like to use the new feature when exporting a SavedModel, simply add the following flag to `tf.saved_model.SaveOptions`:
+
+```python
+tf.saved_model.save(
+  ...,
+  options=tf.saved_model.SaveOptions(experimental_image_format=True)
+)
+```
+
+The Merger has been integrated with `tf.saved_model.load`, so no change needs to be made to SavedModel loading code.
+
+## Chunking Schema
+
+A proto larger than 2GB cannot be serialized. This is a limit of the protobuf implementation that we must work around, which is why we created a proto Splitter/Merger solution. The Splitter takes a proto as input and produces **chunks** and **metadata**. Chunks are parts of a proto that have been split into units of binary data, and can be merged together to form the original proto. Metadata refers to the auxiliary information about where these chunks are extracted from the original proto. This structural information of the proto is contained in the tree-like `ChunkedMessage`. When writing to disk, the metadata takes the form of `ChunkMetadata`, which contains the `ChunkedMessage` as well as information about the chunks' location within the file. When simply splitting the message in memory, only the `ChunkedMessage` is needed. On the Merger side of things, the metadata is used to build the proto back from its disjointed chunks.
+
+`ChunkedMessage` contains an optional `chunk_index`, which references a `chunk` that contains the corresponding message. This message may be further chunked and have one or more of its fields with their own chunks. Therefore, `ChunkedMessage` also contains a list of `ChunkedField`s.
+
+A `ChunkedField` represents a field within a message that has been delegated to its own `chunk`. It contains `field_tag`s that specify where it is located relative to the message `ChunkedField` belongs to. It also contains a `ChunkedMessage`, which allows for a structure that resembles a tree, which is a natural fit for proto metadata.
+
+As an example, consider the following message `A` and its corresponding `ChunkedMessage`:
+
+```proto
+message A {
+  int num = 1;
+  string str = 2;
+  B b = 3;
+}
+
+message B {
+  ...
+}
+```
+
+#### Metadata:
+```proto
+ChunkedMessage {
+  chunk_index: 0
+  chunked_fields: [
+    ChunkedField {
+      field_tag: [b]
+      message: ChunkedMessage {
+        chunk_index: 1
+      }
+    }
+  ]
+}
+```
+
+#### View of memory (deserialized):
+```proto
+chunks [
+  0: A {
+    num: ...
+    str: ...
+  }
+  1: B {
+    ...
+  }
+]
+```
+
+Here, `A`'s `ChunkedMessage` has the optional `chunk_index`, so we see in memory that `chunks[0]` does indeed contain the message `A`. Note that the `A` in `chunks[0]` lacks the `b` field, which has been chunked out. We see this reflected in `A`'s `ChunkedMessage`, whose `chunked_field`s contains the `ChunkedField` that corresponds to this `b` field. The `field_tag`s contain the (very short) path to the `b` field, and the `ChunkedMessage` within the `ChunkedField` references the location of the `chunk` in memory. Indeed, we see the `B` message in memory at `chunks[1]`.
+
+## Field Tag Serialization
+
+A `chunked_field`'s location within the proto is specified by its `field_tag`s.
+
+```proto
+message ChunkedField {
+  repeated FieldIndex field_tag = 1;
+}
+
+message FieldIndex {
+  message MapKey {
+    oneof type {
+      string s = 1;
+      bool boolean = 2;
+      uint32 ui32 = 3;
+      uint64 ui64 = 4;
+      int32 i32 = 5;
+      int64 i64 = 6;
+    }
+  }
+  oneof kind {
+    uint32 field = 1;
+    MapKey map_key = 2;
+    uint64 index = 3;
+  }
+}
+```
+
+Consider the following messages `A`, `B`, and `C`:
+
+```proto
+message A {
+  map<string,B> b = 1;
+}
+
+message B {
+  repeated C c = 1;
+}
+
+message C {
+  BigMessage big = 1;
+}
+
+message BigMessage {
+  ...
+}
+```
+
+Say we were given an `A` proto and wanted to chunk out `big`, since it is quite large. To reference `big`, we use the following path: `A.b["example_string"].c[3].big`. In this case, our list of `field_tag`s would look something like: `[ b, "example_string", c, 3, big ]`. The `field_tag`s for a `chunked_field` (`big`) specify its location relative to the given proto.
+
+These tags represent either a `field`, `map_key`, or `index`, depending on what exactly is being referenced. For example, this allows us to differentiate between `G1 = GraphDef.node.1.attr.value.tensor` and `G2 = GraphDef.node[1].attr["value"].tensor`, even though their lists of `field_tag`s appear to be very similar. `G1`'s `node` field is simply a message containing a field `1`, while `G2`'s `node` field is a repeated message, who's `1`st element is being referenced. Similarly, `G1`'s `attr` field is a message containing a field called `attr`, while `G2`'s `attr` is a map, with the `value` key being referenced. Technically, we could use the proto reflection API to tell whether these ambiguous fields are repeated/map fields or not. However, it is better to be explicit, since it avoids bugs and the extra information makes for a better debugging experience.
+
+## Chunk Extraction and Storage
+
+Proto fields relevant to splitting/merging are classified using their type and occurrence:
+
+ - Field type: **Scalar** or **Message**
+ - Field occurrence: **Singular**, **Repeated**, or **Map**
+
+Other proto field qualifiers like `oneof`, `required`, `optional`, and `packed` do not affect splitting and merging, so they are not taken into account in the implementation.
+
+### Singular Fields
+
+Scalar fields are simply serialized as bytes. Numerical types, such as ints, are serialized in numpy-readable binary. Message fields are also serialized as bytes, once they have been chunked down to <2GB.
+
+### Repeated Fields
+
+When repeated fields are split, they are stored in a chunk that has the same type as the parent of that repeated field. The order of the `chunked_field` for repeated fields is the same order in which the chunks should be merged.
+
+For example, consider the message `A` which contains a repeated field `i`:
+
+```proto
+message A {
+  repeated int i = 1;
+}
+
+A(i=[1, 2, 3, 4, 5])
+```
+
+#### Metadata
+```proto
+ChunkedMessage {
+  chunked_fields: [
+    ChunkedField {
+      field_tag = [],
+      chunk = 0
+    },
+    ChunkedField {
+      field_tag = [],
+      chunk = 1
+    },
+  ]
+}
+```
+
+#### View of memory (deserialized)
+```proto
+chunks [
+  0: A {
+    i=[1, 2]
+  }
+  1: A {
+    i=[3, 4, 5]
+  }
+]
+```
+
+`A`'s `ChunkedMessage` contains two `ChunkedField`s, one for the indices `[1, 2]` and another for the indices `[3, 4, 5]`. The `field_tag`s for both are empty, because the chunks are also of type `A`, and not a field within `A`. During merging, `chunks[0]` must be merged into the in-memory message `A` before `chunks[1]` so that the ordering of the repeated field elements is correct.
+
+### Map Fields
+
+Protobuf maps, like repeated fields, are not a distinct structure within the proto specification. Instead, maps are actually represented by repeated messages with `key` and `value` fields. (This means proto maps aren't really associative containers, but that isn't important here.) Here's an example of a map:
+
+```proto
+message A {
+  map<string, int> my_map = 1;
+}
+A(my_map={"abc": 123, "def": 456})
+```
+
+#### Underlying proto structure:
+```proto
+A: {
+  my_map: {
+    key: "abc"
+    value: 123
+  }
+  my_map: {
+    key: "def"
+    value: 456
+  }
+}
+```
+
+Since maps are really just repeated fields under the hood, we can chunk them the same way we chunk repeated fields:
+
+```proto
+message A {
+  map<int, int> m = 1;
+}
+
+A(i={1:2, 3:4, 5:6})
+```
+
+#### Metadata
+```proto
+ChunkedMessage {
+  chunked_fields: [
+    ChunkedField {
+      field_tag = [],
+      chunk = 0
+    },
+    ChunkedField {
+      field_tag = [],
+      chunk = 1
+    },
+  ]
+}
+```
+
+#### View of memory (deserialized)
+```proto
+chunks [
+  0: A {
+    i={3: 4}
+  }
+  1: A {
+    i={1: 2, 5: 6}
+  }
+]
+```
+
+However, we can also chunk out the values in the map entry directly if we'd like:
+
+```proto
+message A {
+  map<int, B> m = 1;
+}
+
+message B {
+  int i = 1;
+}
+
+A(i={1:B(i=3), 2:B(i=4)})
+```
+
+#### Metadata
+```proto
+ChunkedMessage {
+  chunked_fields: [
+    ChunkedField {
+      field_tag = [m, 3],
+      chunk = 0
+    },
+    ChunkedField {
+      field_tag = [m, 2],
+      chunk = 1
+    },
+  ]
+}
+```
+
+#### View of memory (deserialized)
+```proto
+chunks [
+  0: B {
+    i=3
+  }
+  1: B {
+    i=4
+  }
+]
+```
+
+### Blank Message Compression
+
+In general, we assume the first chunk to be the base message from which all the chunks are extracted (during the split), or the chunk that exists. **However, it's important to note that this isn't required.** If all data is extracted from the user-provided proto into chunks, there is no need for the initial chunk to be the base message. Here's an example with message `A`:
+
+```proto
+message A {
+  B b = 1;
+  C c = 2;
+}
+
+a = A(b=B(...), c=C(...))
+```
+
+Message `a` can be split into chunks `[b, c]` in two ways:
+
+*First chunk is the same as the parent type*
+
+```proto
+chunked_message {
+  chunk_index: 0  // Chunk index is set as the parent message type
+  chunked_fields {  // First field is chunked
+    field_tag { field: 1 }
+    message { chunk_index: 1 }
+  }
+  chunked_fields {  // Second field stored in a separate chunk
+    field_tag { field: 2 }
+    message { chunk_index: 2 }
+  }
+}
+```
+
+#### View of memory (deserialized)
+```proto
+chunks [
+  0: A {...}
+  1: B {...}
+  2: C {...}
+]
+```
+
+*First chunk is not the parent type*
+
+```proto
+chunked_message {
+  // Chunk index is not set in the parent message type
+  chunked_fields {  // First field is chunked
+    field_tag { field: 1 }
+    message { chunk_index: 0 }
+  }
+  chunked_fields {  // Second field stored in a separate chunk
+    field_tag { field: 2 }
+    message { chunk_index: 1 }
+  }
+}
+```
+
+#### View of memory (deserialized)
+```proto
+chunks [
+  0: B {...}
+  1: C {...}
+]
+```
+
+This second method is viable since Message `A` only contains data from fields `b` and `c`. Once `b` and `c` are chunked, there's no other data from `A` to include, so we don't bother creating a chunk for `A`. The merging implementation should not make an assumption on the type of the first chunk, and in this case must create a new (blank) `A` message to merge the `b` and `c` chunks into.
+
+**tldr: A chunked_message may not have a parent chunk to merge its chunked_fields into**
+
+## Creating a Splitter
+
+Now that we've covered the format used by the Splitters/Merger, we can work on implementing our own Splitter. By now you can understand why each proto requires its own bespoke Splitter, since automatic splitting wouldn't take advantage of the knowledge we have as proto designers of bottlenecks and opportunities for optimization. So, let's walk through the process of creating a Splitter for our message `ModelConfig`:
+
+```proto
+enum ActivationFunction {
+  RELU = 0;
+  SIGMOID = 1;
+  TANH = 2;
+}
+
+message Layer {
+  string name = 1;
+  int32 num_units = 2;
+  ActivationFunction activation_function = 3;
+}
+
+message ModelConfig {
+  string model_name = 1;
+  int32 input_shape = 2;
+  repeated Layer hidden_layers = 3;
+  int32 output_units = 4;
+  ActivationFunction output_activation = 5;
+  map<string, float> hyperparameters = 6;
+}
+```
+
+To create a `ModelConfig` Splitter, we have to decide what exactly is being split. As the designers of `ModelConfig`, we know that the `hidden_layers` tend to be quite large, so that makes the `Layer`s messages good candidates to split out into their own chunks. For the sake of example, we're also going to split out the `hyperparameters` field.
+
+To create a Splitter, we must subclass the `ComposableSplitter` class and override its `build_chunks` method. If we wanted to store state in a Splitter, we could also override the `__init__` method, but it isn't required. In our example this would be enough to split and chunk out the fields we settled on (`hidden_layers` and `hyperparameters`), but we'll also create a Splitter for the `Layer` message to showcase Splitter composition.
+
+```python
+class ModelConfigSplitter(ComposableSplitter):
+  def build_chunks(self):
+    for k, v in self._proto.hyperparameters:
+      self.add_chunk(bytes(str(v), "utf-8"), ["hyperparameters", k])
+
+    for i, layer in enumerate(self._proto.hidden_layers):
+      LayerSplitter(
+        layer,
+        parent_splitter=self,
+        fields_in_parent=["hidden_layers", i]
+      ).build_chunks()
+
+class LayerSplitter(ComposableSplitter):
+  def build_chunks(self):
+    self.add_chunk(self._proto, [])
+
+ModelConfigSplitter(
+  proto=ModelConfig(...)
+)
+```
+
+`build_chunks` generates chunks from `self._proto`, then for each chunk, calls `add_chunk` to add it to `self._chunks` and update `self._chunked_message`. `ModelConfigSplitter` does this once for `hyperparameters`, by simply converting the float value to a string and then to bytes. The Splitter does it again for `hidden_layers`, which get chunked by a dedicated `LayerSplitter` class. `LayerSplitter` doesn't actually do any chunking, but is here to showcase the ability to have a hierarchy of Splitters.
+
+## Merging
+
+There are two ways of merging a chunked proto using the provided Merger:
+
+ - `Merger::Read()`, merges directly into a user-provided merged_message from a .cpb file on disk
+ - `Merger::Merge()`, requires that the chunks and chunked metadata be stored in memory
+
+`Merge()` should be called at runtime with the C++ Splitter, and allows one to skip any unnecessary disk reads/writes. `Read()` is therefore more holistic, handling both file IO and merging, so we'll consider its implementation below. The provided Merger is independent of any Splitter or protobuf, so developers will not have to write their own in the vast majority of cases.
+
+### Riegeli
+
+Since chunked protos use the riegeli file format, we use the riegeli api for file IO. The `riegeli::RecordReader` makes it easy to `Seek()` to a position in the file and `ReadRecord()` at that location.
+
+### Reflection
+
+We also make use of the protobuf reflection api to add and modify fields in `merged_message` using `FieldDescriptor`s.
+
+### ChunkedMetadata
+
+But to understand what should be read and where to read it from, we need the `ChunkedMetadata`. The metadata is always stored in the last chunk of the chunked proto, so we simply read that record to begin the merging process. Within the `ChunkedMetadata`, the sequence of `ChunkInfo` tells us where in the chunked proto to find the chunk we're looking for. And the `ChunkedMessage` contains a tree of metadata that we can use to reconstruct the desired proto.
+
+### Field Processing
+
+Starting at the root `ChunkedMessage`, we first check to see if it references a chunk by specifying a `chunk_index`. If so, we need to merge that chunk into the target proto (let's call it `A`) before processing each of its `chunked_field`s. If there is no `chunk_index`, then `A` only contains fields that have been chunked out. Before merging in the `chunked_field`s, they must be sorted by depth and index. For example, we need to merge in `GraphDef.library` before `GraphDef.library.function[0]`, which needs to be merged in before `GraphDef.library.function[1]`. We must merge in the `library` field first so that the `library.function`s have some place to be merged into, and the `0`th `function` must be merged before the `1`st `function` to maintain the proper ordering. Now we're ready to merge in the `chunked_field`s.
+
+For each `ChunkedField` in a `ChunkedMessage`:
+
+1. Read in the `chunk` specified by the `chunks_info[chunked_field.message.chunk_index]`
+2. If the `chunked_field` has no `field_tag`s, then it does not reference a field within the parent message, but rather part of the parent message itself. For example, consider the following message and its corresponding `chunked_message`:
+   ```proto
+   message A {
+     ...
+   }
+   
+   chunked_message = {
+     chunked_fields {  // empty field_tag, belongs to the parent chunked_message
+       field_tag { }
+       message { chunk_index: 0}
+     }
+     chunked_fields {  // also belongs to the parent
+    chunk
+       field_tag { }
+       message { chunk_index: 1}
+     }
+   }
+   ```
+   In this case, a message `A` has been split into multiple chunks (here `A1` and `A2`, but hypothetically up to `An`), rather than splitting its fields into their own chunks. Splitting a message into chunks directly or splitting a message's fields into chunks are simply two different approaches that we offer in our api. So, the `chunk` should be merged directly into the parent message (`A`), and we skip the remaining steps to move on to the next `chunked_field`.
+3. Navigate the `merged_message` using the `field_tag`s, until reaching the target field. Fields may need to be constructed along the way if they were not kept during the splitting process (see [Blank Message Compression above](#blank_message_compression)).
+4. If the field is not a message, it is a primitive data type like bool or int, so we simply convert the `chunk` string to the appropriate type and set the field using reflection. If it is a message, then we recursively process it using its corresponding `ChunkedMessage`.
+
+When the recursive process is complete, the `chunk`s have been successfully merged into the `merged_message`, so it's ready to be used in your program.
+
+## Putting It All Together
+
+Now that we've covered the entire splitting and merging process, let's go over an end-to-end example. We'll use the `ModelConfigSplitter` class we created in the [Creating a Splitter](#creating_a_splitter) section above. To write our proto to disk, we simply call `Splitter.write()`:
+
+```python
+my_proto = ModelConfig(...)
+export_dir = "..."
+my_splitter = ModelConfigSplitter(my_proto)
+my_splitter.write(export_dir)
+```
+
+And in C++, we can use the Merger to read in our chunked proto:
+
+```c++
+ModelConfig my_proto;
+string export_dir = "...";
+Merger::Read(export_dir, &my_proto);
+```
+
+If we'd instead like to split and merge our proto directly in memory, we'd need `ModelConfigSplitter` to be a C++ class, but the process is very similar:
+
+```c++
+class ModelConfigSplitter : public ComposableSplitter {
+  ...
+};
+
+ModelConfig my_proto{...};
+string export_dir = "...";
+ModelConfigSplitter my_splitter(my_proto);
+
+// std::pair<std::vector<MessageBytes>*, ::proto_splitter::ChunkedMessage*>
+auto[chunks, chunked_message] = my_splitter.Split();
+
+// chunks, chunked_message are processed
+
+ModelConfig my_new_proto;
+Merger::Merge(chunks, chunked_message, &my_new_proto);
+```

tensorflow/tools/proto_splitter/merge.cc

@@ -0,0 +1,440 @@
+/* Copyright 2023 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/tools/proto_splitter/merge.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "riegeli/base/object.h"  // from @riegeli
+#include "riegeli/bytes/fd_reader.h"  // from @riegeli
+#include "riegeli/records/record_reader.h"  // from @riegeli
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system_helper.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tsl/platform/errors.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow::tools::proto_splitter {
+
+using ::tensorflow::proto_splitter::ChunkedField;
+using ::tensorflow::proto_splitter::ChunkedMessage;
+using ::tensorflow::proto_splitter::ChunkInfo;
+using ::tensorflow::proto_splitter::ChunkMetadata;
+using ::tensorflow::proto_splitter::FieldIndex;
+using tools::proto_splitter::GetChunkMetadata;
+using tools::proto_splitter::GetRiegeliReader;
+using tools::proto_splitter::OnlyContainsPb;
+using tsl::protobuf::FieldDescriptor;
+using tsl::protobuf::Message;
+using tsl::protobuf::Reflection;
+
+absl::Status Merger::Merge(const std::vector<std::unique_ptr<Message>>& chunks,
+                           const ChunkedMessage& chunked_message,
+                           Message* merged_message) {
+  riegeli::RecordReader<riegeli::FdReader<>> null_reader{riegeli::kClosed};
+
+  if (chunked_message.has_chunk_index()) {
+    // Chunks referenced by fields should be merged into the parent chunk.
+    merged_message->MergeFrom(*chunks[chunked_message.chunk_index()].get());
+  }
+
+  // Use each chunked_field within the chunked_message to merge its
+  // corresponding chunk into merged_message.
+  for (const auto& chunked_field : chunked_message.chunked_fields()) {
+    absl::Status s = ProcessField(chunked_field, merged_message, {}, chunks,
+                                  null_reader, MergerOp::MERGE);
+    if (!s.ok()) return s;
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status Merger::Read(std::string prefix, Message* merged_message) {
+  uint64_t start_time = Env::Default()->NowMicros();
+
+  TF_ASSIGN_OR_RETURN(bool only_contains_pb, OnlyContainsPb(prefix));
+  if (only_contains_pb) {
+    return ReadPb(absl::StrCat(prefix, ".pb"), merged_message);
+  }
+
+  // Create riegeli reader for file.cpb
+  TF_ASSIGN_OR_RETURN(auto reader,
+                      GetRiegeliReader(absl::StrCat(prefix, ".cpb")));
+
+  auto read_metadata = GetChunkMetadata(reader);
+  if (!read_metadata.ok()) {
+    reader.Close();
+    return absl::FailedPreconditionError(
+        absl::StrCat("Couldn't read ChunkMetadata from chunked proto.\n",
+                     read_metadata.status().ToString()));
+  }
+  ChunkMetadata chunk_metadata = read_metadata.value();
+
+  std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
+      chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
+
+  // Read the remaining chunks.
+  absl::Status s =
+      ReadFields(chunk_metadata.message(), reader, chunks_info, merged_message);
+
+  reader.Close();
+
+  uint64_t end_time = Env::Default()->NowMicros();
+  LOG(INFO) << "Finished reading and merging chunked proto, took "
+            << HumanReadableDuration(end_time - start_time) << ".";
+  return s;
+}
+
+absl::Status Merger::ReadPartial(absl::string_view prefix,
+                                 const ChunkMetadata& chunk_metadata,
+                                 Message* merged_message) {
+  uint64_t start_time = Env::Default()->NowMicros();
+
+  TF_ASSIGN_OR_RETURN(bool only_contains_pb, OnlyContainsPb(prefix));
+  if (only_contains_pb) {
+    return absl::FailedPreconditionError(
+        absl::StrCat("Attempting to read part of a chunked proto .cpb file, "
+                     "but only found a regular proto: ",
+                     prefix, ".pb"));
+  }
+
+  // Create riegeli reader for file.cpb
+  TF_ASSIGN_OR_RETURN(auto reader,
+                      GetRiegeliReader(absl::StrCat(prefix, ".cpb")));
+
+  std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
+      chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
+
+  // Read the remaining chunks.
+  absl::Status s =
+      ReadFields(chunk_metadata.message(), reader, chunks_info, merged_message);
+
+  reader.Close();
+
+  uint64_t end_time = Env::Default()->NowMicros();
+  LOG(INFO) << "Finished reading and merging chunked proto, took "
+            << HumanReadableDuration(end_time - start_time) << ".";
+  return s;
+}
+
+absl::Status Merger::ReadChunkedFromString(
+    absl::string_view data, tsl::protobuf::Message* merged_message) {
+  uint64_t start_time = Env::Default()->NowMicros();
+
+  // Create riegeli string reader.
+  TF_ASSIGN_OR_RETURN(auto reader, GetRiegeliStringReader(data));
+
+  auto read_metadata = GetChunkMetadata(reader);
+  if (!read_metadata.ok()) {
+    reader.Close();
+    return absl::FailedPreconditionError(
+        absl::StrCat("Couldn't read ChunkMetadata from chunked proto.\n",
+                     read_metadata.status().ToString()));
+  }
+  ChunkMetadata chunk_metadata = read_metadata.value();
+
+  std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
+      chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
+
+  // Read the remaining chunks.
+  absl::Status s =
+      ReadFields(chunk_metadata.message(), reader, chunks_info, merged_message);
+
+  reader.Close();
+
+  uint64_t end_time = Env::Default()->NowMicros();
+  LOG(INFO) << "Finished reading and merging chunked proto, took "
+            << HumanReadableDuration(end_time - start_time) << ".";
+  return s;
+}
+
+absl::Status Merger::ReadPb(const std::string& pb_file,
+                            Message* merged_message) {
+  uint64_t start_time = Env::Default()->NowMicros();
+  TF_ASSIGN_OR_RETURN(bool file_exists,
+                      internal::FileExists(Env::Default(), pb_file));
+  if (!file_exists)
+    return absl::NotFoundError(absl::StrCat("File not found: ", pb_file));
+  LOG(INFO) << "Reading binary proto from " << pb_file;
+  auto ret = ReadBinaryProto(Env::Default(), pb_file, merged_message);
+  uint64_t end_time = Env::Default()->NowMicros();
+  LOG(INFO) << "Finished reading binary proto, took "
+            << HumanReadableDuration(end_time - start_time) << ".";
+  return ret;
+}
+
+namespace {
+
+bool CompareMapKey(const FieldIndex::MapKey& key1,
+                   const FieldIndex::MapKey& key2) {
+  if (key1.type_case() != key2.type_case()) {
+    return key1.type_case() < key2.type_case();
+  }
+  switch (key1.type_case()) {
+    case FieldIndex::MapKey::TypeCase::kS:
+      return key1.s() < key2.s();
+    case FieldIndex::MapKey::TypeCase::kBoolean:
+      return key1.boolean() < key2.boolean();
+    case FieldIndex::MapKey::TypeCase::kUi32:
+      return key1.ui32() < key2.ui32();
+    case FieldIndex::MapKey::TypeCase::kUi64:
+      return key1.ui64() < key2.ui64();
+    case FieldIndex::MapKey::TypeCase::kI32:
+      return key1.i32() < key2.i32();
+    case FieldIndex::MapKey::TypeCase::kI64:
+      return key1.i64() < key2.i64();
+    default:
+      return false;
+  }
+}
+
+bool CompareFieldIndex(const FieldIndex& tag1, const FieldIndex& tag2) {
+  if (tag1.kind_case() != tag2.kind_case()) {
+    return tag1.kind_case() < tag2.kind_case();
+  }
+  switch (tag1.kind_case()) {
+    case FieldIndex::KindCase::kField:
+      return tag1.field() < tag2.field();
+    case FieldIndex::KindCase::kIndex:
+      return tag1.index() < tag2.index();
+    case FieldIndex::KindCase::kMapKey:
+      return CompareMapKey(tag1.map_key(), tag2.map_key());
+    default:
+      return false;
+  }
+}
+
+}  // namespace
+
+template <typename RecordReader>
+absl::Status Merger::ReadFields(const ChunkedMessage& chunked_message,
+                                RecordReader& reader,
+                                const std::vector<ChunkInfo>& chunks_info,
+                                tsl::protobuf::Message* merged_message) {
+  if (chunked_message.has_chunk_index()) {
+    // Chunks referenced by fields should be merged into the parent chunk.
+    TF_ASSIGN_OR_RETURN(
+        std::string chunk,
+        ReadChunk(reader, chunks_info[chunked_message.chunk_index()]));
+    if (!merged_message->MergeFromString(chunk)) {
+      return absl::FailedPreconditionError(
+          "Couldn't merge chunk into message.");
+    }
+  }
+
+  // Sort the chunked_fields by depth and index.
+  // For example, this ensures that GraphDef.library is merged before
+  //   GraphDef.library.function[0], which will be merged before
+  //   GraphDef.library.function[1].
+  std::vector<ChunkedField> chunked_fields(
+      chunked_message.chunked_fields().begin(),
+      chunked_message.chunked_fields().end());
+  std::sort(
+      chunked_fields.begin(), chunked_fields.end(),
+      // Ensure proper strict weak ordering.
+      [](const ChunkedField& cf1, const ChunkedField& cf2) {
+        int tag_depth =
+            std::min(cf1.field_tag().size(), cf2.field_tag().size());
+        for (int depth = 0; depth < tag_depth; ++depth) {
+          const FieldIndex& tag1 = cf1.field_tag()[depth];
+          const FieldIndex& tag2 = cf2.field_tag()[depth];
+          bool is_equiv = false;
+          if (tag1.kind_case() == tag2.kind_case()) {
+            if (tag1.kind_case() == FieldIndex::KindCase::kField) {
+              is_equiv = (tag1.field() == tag2.field());
+            } else if (tag1.kind_case() == FieldIndex::KindCase::kIndex) {
+              is_equiv = (tag1.index() == tag2.index());
+            } else if (tag1.kind_case() == FieldIndex::KindCase::kMapKey) {
+              is_equiv = !CompareMapKey(tag1.map_key(), tag2.map_key()) &&
+                         !CompareMapKey(tag2.map_key(), tag1.map_key());
+            } else {
+              is_equiv = true;  // KIND_NOT_SET
+            }
+          }
+          if (!is_equiv) {
+            return CompareFieldIndex(tag1, tag2);
+          }
+        }
+        if (cf1.field_tag().size() == cf2.field_tag().size()) {
+          // If the fields are the same, merge the earlier chunks first.
+          return cf1.message().chunk_index() < cf2.message().chunk_index();
+        }
+        return cf1.field_tag().size() < cf2.field_tag().size();
+      });
+
+  // Use each chunked_field within the chunked_message to merge its
+  // corresponding chunk into merged_message.
+  for (const auto& chunked_field : chunked_fields) {
+    absl::Status s = ProcessField(chunked_field, merged_message, chunks_info,
+                                  {}, reader, MergerOp::READ);
+    if (!s.ok()) return s;
+  }
+  return absl::OkStatus();
+}
+
+template <typename RecordReader>
+absl::Status Merger::ProcessField(
+    const ChunkedField& chunked_field, Message* merged_message,
+    const std::vector<ChunkInfo>& chunks_info,
+    const std::vector<std::unique_ptr<Message>>& chunks, RecordReader& reader,
+    MergerOp op) {
+  std::string chunk;
+  switch (op) {
+    case MergerOp::READ: {
+      TF_ASSIGN_OR_RETURN(
+          chunk, ReadChunk(reader,
+                           chunks_info[chunked_field.message().chunk_index()]));
+      break;
+    }
+    case MergerOp::MERGE: {
+      chunk =
+          chunks[chunked_field.message().chunk_index()]->SerializeAsString();
+      break;
+    }
+  }
+
+  if (chunked_field.field_tag().empty()) {
+    // Chunk is not a field within the parent, but instead a portion of the
+    // parent itself. Needs to be concatenated.
+    merged_message->MergeFromString(chunk);
+    return absl::OkStatus();
+  }
+
+  uint64_t field_index;
+  Message* curr_message = merged_message;
+  // Use field tags to navigate the merged_message, constructing necessary
+  // fields along the way.
+  TF_ASSIGN_OR_RETURN(const std::vector<Field> fields,
+                      GetFieldTypes(chunked_field.field_tag()));
+  const FieldDescriptor* field_desc = nullptr;
+  for (const auto& field : fields) {
+    merged_message = curr_message;
+    field_desc = merged_message->GetDescriptor()->FindFieldByNumber(
+        std::get<int>(field.first));
+
+    if (field_desc == nullptr) {
+      return absl::FailedPreconditionError(
+          absl::StrCat("Field with number ", std::get<int>(field.first),
+                       " not found in message descriptor.",
+                       merged_message->GetDescriptor()->full_name(), ".\n"));
+    }
+
+    auto res = GetMutableField(merged_message, field);
+    if (!res.ok()) {
+      if (!absl::IsNotFound(res.status())) return res.status();
+      // Add missing field.
+      if (field_desc->is_map()) {
+        TF_RETURN_IF_ERROR(
+            AddMapEntry(curr_message, field_desc, field.second.value()));
+        res = GetMutableField(curr_message, field);
+      } else {
+        curr_message->GetReflection()->AddMessage(curr_message, field_desc);
+        res = GetMutableField(curr_message, field);
+      }
+    }
+
+    auto [parent, mutable_field, mutable_field_index] = res.value();
+    if (mutable_field->is_repeated() && mutable_field_index != -1) {
+      field_index = mutable_field_index;
+      // Update merged_message to repeated element.
+      curr_message = parent->GetReflection()->MutableRepeatedMessage(
+          parent, mutable_field, std::max(0, mutable_field_index));
+      if (mutable_field->is_map()) {
+        // messages of map type have the value at field #2
+        field_desc = mutable_field->message_type()->FindFieldByNumber(2);
+        merged_message = curr_message;
+        curr_message = curr_message->GetReflection()->MutableMessage(
+            curr_message, field_desc);
+      }
+    } else if (mutable_field->type() == FieldDescriptor::Type::TYPE_MESSAGE) {
+      // Update merged_message.
+      curr_message =
+          parent->GetReflection()->MutableMessage(parent, mutable_field);
+    }
+  }
+  // merged_message now points to the Message whose field (described by
+  // field_desc) will be added/set.
+
+  const Reflection* reflection = merged_message->GetReflection();
+  if (field_desc->is_repeated()) {
+    // field may contain multiple elements
+    auto message_callback = [&reflection, &merged_message, &field_index, &op,
+                             &chunks, &chunked_field, &reader, &chunks_info,
+                             &field_desc]() -> absl::Status {
+      for (int _ = reflection->FieldSize(*merged_message, field_desc);
+           _ <= field_index; _++) {
+        reflection->AddMessage(merged_message, field_desc);
+      }
+      switch (op) {
+        case MergerOp::MERGE:
+          TF_RETURN_IF_ERROR(
+              Merge(chunks, chunked_field.message(),
+                    reflection->MutableRepeatedMessage(
+                        merged_message, field_desc, field_index)));
+          break;
+        case MergerOp::READ:
+          TF_RETURN_IF_ERROR(
+              ReadFields(chunked_field.message(), reader, chunks_info,
+                         reflection->MutableRepeatedMessage(
+                             merged_message, field_desc, field_index)));
+          break;
+        default:
+          return absl::InternalError("Encountered unknown MergerOp.");
+      }
+      return absl::OkStatus();
+    };
+    TF_RETURN_IF_ERROR(SetRepeatedFieldElement(
+        merged_message, field_desc, field_index, chunk, message_callback));
+  } else {
+    // regular field
+    auto message_callback = [&reflection, &merged_message, &op, &chunks,
+                             &chunked_field, &reader, &chunks_info,
+                             &field_desc]() -> absl::Status {
+      switch (op) {
+        case MergerOp::MERGE:
+          TF_RETURN_IF_ERROR(
+              Merge(chunks, chunked_field.message(),
+                    reflection->MutableMessage(merged_message, field_desc)));
+          break;
+        case MergerOp::READ:
+          TF_RETURN_IF_ERROR(ReadFields(
+              chunked_field.message(), reader, chunks_info,
+              reflection->MutableMessage(merged_message, field_desc)));
+          break;
+        default:
+          return absl::InternalError("Encountered unknown MergerOp.");
+      }
+      return absl::OkStatus();
+    };
+    TF_RETURN_IF_ERROR(
+        SetFieldElement(merged_message, field_desc, chunk, message_callback));
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow::tools::proto_splitter

tensorflow/tools/proto_splitter/merge.h

@@ -0,0 +1,97 @@
+/* Copyright 2023 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_TOOLS_PROTO_SPLITTER_MERGE_H_
+#define TENSORFLOW_TOOLS_PROTO_SPLITTER_MERGE_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow::tools::proto_splitter {
+
+class Merger {
+ private:
+  enum MergerOp { MERGE, READ };
+
+ public:
+  // Merges the provided `chunks` into `merged_message` using `chunked_message`.
+  // Example usage:
+  //   std::vector<tsl::protobuf::Message> chunks = GetMyChunks();
+  //   ::proto_splitter::ChunkedMessage chunked_message = GetMyChunkedMessage();
+  //   my_project::MyProto my_proto;
+  //   Merger::Merge(chunks, chunked_message, &my_proto);
+  // TODO(b/282775853): Integrate Splitter return type with Merge input type
+  static absl::Status Merge(
+      const std::vector<std::unique_ptr<tsl::protobuf::Message>>& chunks,
+      const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
+      tsl::protobuf::Message* merged_message);
+
+  // Reads a TF SavedModel chunked protobuf from `prefix` (must be .pb or .cpb)
+  // into `merged_message`. The proto format of `merged_message` must match the
+  // format of the proto written to `prefix`.
+  // Example usage:
+  //   my_project::MyProto my_proto;
+  //   Merger::Read("path/to/saved_model", &my_proto);
+  static absl::Status Read(std::string prefix,
+                           tsl::protobuf::Message* merged_message);
+
+  // Like `Merger::Read`, but only reads what's specified in `chunk_metadata`.
+  static absl::Status ReadPartial(
+      absl::string_view prefix,
+      const ::tensorflow::proto_splitter::ChunkMetadata& chunk_metadata,
+      tsl::protobuf::Message* merged_message);
+
+  // Deserializes a chunked protobuf from a string into `merged_message`.
+  static absl::Status ReadChunkedFromString(
+      absl::string_view data, tsl::protobuf::Message* merged_message);
+
+ private:
+  // Reads a normal saved_model.pb proto in.
+  static absl::Status ReadPb(const std::string& pb_file,
+                             tsl::protobuf::Message* merged_message);
+
+  // Uses metadata contained in `chunked_message` to fill `merged_message` with
+  // data accessed by the `reader` using `chunks_info`.
+  template <typename RecordReader>
+  static absl::Status ReadFields(
+      const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
+      RecordReader& reader,
+      const std::vector<::tensorflow::proto_splitter::ChunkInfo>&
+          chunks_info,  // TODO(adamcogdell): this can just be a
+                        // RepeatedPtrField
+      tsl::protobuf::Message* merged_message);
+
+  // Processes a single `chunked_field` within a `chunked_message`. If the field
+  // itself is a `chunked_message` that contains additional `chunked_fields`,
+  // either MergeFields or ReadFields is called to recursively (depending on the
+  // value of `op`) to add those fields to `merged_message`. Otherwise, the
+  // field is simply added to `merged_message` using reflection.
+  template <typename RecordReader>
+  static absl::Status ProcessField(
+      const ::tensorflow::proto_splitter::ChunkedField& chunked_field,
+      tsl::protobuf::Message* merged_message,
+      const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info,
+      const std::vector<std::unique_ptr<tsl::protobuf::Message>>& chunks,
+      RecordReader& reader, MergerOp op);
+};
+
+}  // namespace tensorflow::tools::proto_splitter
+
+#endif  // TENSORFLOW_TOOLS_PROTO_SPLITTER_MERGE_H_

tensorflow/tools/proto_splitter/merge_test.cc

@@ -0,0 +1,317 @@
+/* Copyright 2023 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/tools/proto_splitter/merge.h"
+
+#include <array>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/env.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/tools/proto_splitter/cc/test_util.h"
+#include "tensorflow/tools/proto_splitter/cc/util.h"
+#include "tensorflow/tools/proto_splitter/chunk.pb.h"
+#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
+#include "tsl/platform/env.h"
+#include "tsl/platform/protobuf.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow::tools::proto_splitter {
+
+namespace {
+
+inline constexpr std::array kDFSplitTreeChunks = {
+    "val: \"0\"",       "val: \"010\"",     "val: \"01020\"",
+    "val: \"0102030\"", "val: \"0102031\"", "val: \"0102032\"",
+    "val: \"01021\"",   "val: \"0102130\"", "val: \"0102131\"",
+    "val: \"0102132\""};
+
+inline constexpr std::array kBFSplitTreeChunks = {
+    "val: \"0\"",       "val: \"010\"",     "val: \"01020\"",
+    "val: \"01021\"",   "val: \"0102030\"", "val: \"0102031\"",
+    "val: \"0102032\"", "val: \"0102130\"", "val: \"0102131\"",
+    "val: \"0102132\""};
+
+TEST(MergeTest, TestReadRiegeliTreeDepthFirst) {
+  // TODO(b/282779639): Use test data.
+  const std::string cpb_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "df-split-tree");
+  ::tensorflow::proto_splitter_testdata::StringNode merged_tree;
+  TF_ASSERT_OK(Merger::Read(cpb_path, &merged_tree));
+
+  const std::string pbtxt_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-tree");
+  ::tensorflow::proto_splitter_testdata::StringNode test_proto;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(pbtxt_path, ".pbtxt"), &test_proto));
+
+  ASSERT_THAT(merged_tree, EqualsProto(test_proto));
+}
+
+TEST(MergeTest, TestReadRiegeliTreeBreadthFirst) {
+  const std::string cpb_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "bf-split-tree");
+  ::tensorflow::proto_splitter_testdata::StringNode merged_tree;
+  TF_ASSERT_OK(Merger::Read(cpb_path, &merged_tree));
+
+  const std::string pbtxt_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-tree");
+
+  ::tensorflow::proto_splitter_testdata::StringNode test_proto;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(pbtxt_path, ".pbtxt"), &test_proto));
+
+  ASSERT_THAT(merged_tree, EqualsProto(test_proto));
+}
+
+TEST(MergeTest, TestMergeTreeChunksDepthFirst) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "df-split-tree");
+  std::vector<std::unique_ptr<::tsl::protobuf::Message>> chunks;
+  for (const auto& chunk : kDFSplitTreeChunks) {
+    ::tensorflow::proto_splitter_testdata::StringNode string_node;
+    ::tsl::protobuf::TextFormat::ParseFromString(chunk, &string_node);
+    std::unique_ptr<::tsl::protobuf::Message> node =
+        std::make_unique<::tensorflow::proto_splitter_testdata::StringNode>(
+            string_node);
+    chunks.push_back(std::move(node));
+  }
+
+  std::string split_tree_metadata;
+  TF_ASSERT_OK(tsl::ReadFileToString(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &split_tree_metadata));
+  ::tensorflow::proto_splitter::ChunkedMessage chunked_message;
+  ::tsl::protobuf::TextFormat::ParseFromString(split_tree_metadata,
+                                               &chunked_message);
+
+  ::tensorflow::proto_splitter_testdata::StringNode merged_tree;
+  TF_ASSERT_OK(Merger::Merge(chunks, chunked_message, &merged_tree));
+
+  const std::string pbtxt_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-tree");
+
+  ::tensorflow::proto_splitter_testdata::StringNode test_proto;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(pbtxt_path, ".pbtxt"), &test_proto));
+
+  ASSERT_THAT(merged_tree, EqualsProto(test_proto));
+}
+
+TEST(MergeTest, TestMergeTreeChunksBreadthFirst) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "bf-split-tree");
+  std::vector<std::unique_ptr<::tsl::protobuf::Message>> chunks;
+  for (const auto& chunk : kBFSplitTreeChunks) {
+    ::tensorflow::proto_splitter_testdata::StringNode string_node;
+    ::tsl::protobuf::TextFormat::ParseFromString(chunk, &string_node);
+    std::unique_ptr<::tsl::protobuf::Message> node =
+        std::make_unique<::tensorflow::proto_splitter_testdata::StringNode>(
+            string_node);
+    chunks.push_back(std::move(node));
+  }
+
+  std::string split_tree_metadata;
+  TF_ASSERT_OK(tsl::ReadFileToString(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &split_tree_metadata));
+  ::tensorflow::proto_splitter::ChunkedMessage chunked_message;
+  ::tsl::protobuf::TextFormat::ParseFromString(split_tree_metadata,
+                                               &chunked_message);
+
+  ::tensorflow::proto_splitter_testdata::StringNode merged_tree;
+  TF_ASSERT_OK(Merger::Merge(chunks, chunked_message, &merged_tree));
+
+  const std::string pbtxt_path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-tree");
+
+  ::tensorflow::proto_splitter_testdata::StringNode test_proto;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(pbtxt_path, ".pbtxt"), &test_proto));
+
+  ASSERT_THAT(merged_tree, EqualsProto(test_proto));
+}
+
+TEST(MergeTest, TestReadGraphDefLotsNodes) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-lots-nodes");
+  GraphDef merged_graph_def;
+  TF_ASSERT_OK(Merger::Read(path, &merged_graph_def));
+
+  GraphDef test_graph_def;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &test_graph_def));
+
+  ASSERT_THAT(merged_graph_def, EqualsProto(test_graph_def));
+}
+
+TEST(MergeTest, TestReadGraphDefLargeNodes) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-large-nodes");
+  GraphDef merged_graph_def;
+  TF_ASSERT_OK(Merger::Read(path, &merged_graph_def));
+
+  GraphDef test_graph_def;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &test_graph_def));
+
+  ASSERT_THAT(merged_graph_def, EqualsProto(test_graph_def));
+}
+
+TEST(MergeTest, TestReadGraphDefLargeConstant) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-large-constant");
+  GraphDef merged_graph_def;
+  TF_ASSERT_OK(Merger::Read(path, &merged_graph_def));
+
+  GraphDef test_graph_def;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &test_graph_def));
+
+  ASSERT_THAT(merged_graph_def, EqualsProto(test_graph_def));
+}
+
+TEST(MergeTest, TestReadManyField) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "many-field");
+  ::tensorflow::proto_splitter_testdata::ManyFields merged_many_field;
+  TF_ASSERT_OK(Merger::Read(path, &merged_many_field));
+
+  ::tensorflow::proto_splitter_testdata::ManyFields test_many_field;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &test_many_field));
+
+  ASSERT_THAT(merged_many_field, EqualsProto(test_many_field));
+}
+
+TEST(MergeTest, TestReadSavedModel) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "split-standard");
+  SavedModel merged_saved_model;
+  TF_ASSERT_OK(Merger::Read(path, &merged_saved_model));
+
+  SavedModel test_saved_model;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &test_saved_model));
+
+  ASSERT_THAT(merged_saved_model, EqualsProto(test_saved_model));
+}
+
+TEST(MergeTest, TestReadChunkedModel) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
+                   "chunked_saved_model/chunked_model/saved_model");
+  SavedModel merged_saved_model;
+  TF_ASSERT_OK(Merger::Read(path, &merged_saved_model));
+
+  SavedModel test_saved_model;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &test_saved_model));
+
+  ASSERT_THAT(merged_saved_model, EqualsProto(test_saved_model));
+}
+
+TEST(MergeTest, TestReadPartial) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(),
+                   "tools/proto_splitter/testdata", "many-field");
+  TF_ASSERT_OK_AND_ASSIGN(auto reader, tools::proto_splitter::GetRiegeliReader(
+                                           absl::StrCat(path, ".cpb")));
+
+  auto read_metadata = GetChunkMetadata(reader);
+  if (!read_metadata.ok()) {
+    reader.Close();
+    TF_ASSERT_OK(read_metadata.status());
+  }
+  ::tensorflow::proto_splitter::ChunkMetadata chunk_metadata =
+      read_metadata.value();
+  ::tensorflow::proto_splitter::ChunkMetadata partial_chunk_metadata;
+  partial_chunk_metadata.mutable_chunks()->CopyFrom(chunk_metadata.chunks());
+  partial_chunk_metadata.mutable_message()->set_chunk_index(
+      chunk_metadata.message().chunk_index());
+
+  proto_splitter_testdata::ManyFields merged_many_fields;
+  TF_ASSERT_OK(
+      Merger::ReadPartial(path, partial_chunk_metadata, &merged_many_fields));
+  ASSERT_THAT(merged_many_fields, EqualsProto(R"pb(
+                map_field_int64 { key: -1345 value: "map_value_-1345" }
+              )pb"));
+}
+
+TEST(MergeTest, TestReadChunkedFromString) {
+  const std::string path =
+      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
+                   "chunked_saved_model/chunked_model/saved_model");
+  std::string data;
+  TF_ASSERT_OK(tsl::ReadFileToString(tsl::Env::Default(),
+                                     absl::StrCat(path, ".cpb"), &data));
+
+  SavedModel merged_saved_model;
+  TF_ASSERT_OK(Merger::ReadChunkedFromString(data, &merged_saved_model));
+
+  SavedModel test_saved_model;
+  TF_ASSERT_OK(tsl::ReadTextProto(
+      tsl::Env::Default(), absl::StrCat(path, ".pbtxt"), &test_saved_model));
+
+  ASSERT_THAT(merged_saved_model, EqualsProto(test_saved_model));
+}
+
+TEST(MergeTest, TestProcessFieldReturnsErrorOnInvalidFieldNumber) {
+  ::tensorflow::proto_splitter::ChunkedMessage chunked_message;
+
+  auto* chunk_field = chunked_message.add_chunked_fields();
+
+  auto* tag = chunk_field->add_field_tag();
+  tag->set_field(99999);
+
+  chunk_field->mutable_message()->set_chunk_index(0);
+
+  std::vector<std::unique_ptr<tsl::protobuf::Message>> chunks;
+  chunks.push_back(
+      std::make_unique<::tensorflow::proto_splitter_testdata::ManyFields>());
+
+  ::tensorflow::proto_splitter_testdata::ManyFields merged;
+  absl::Status status = Merger::Merge(chunks, chunked_message, &merged);
+
+  EXPECT_FALSE(status.ok());
+  EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition);
+  EXPECT_THAT(std::string(status.message()),
+              ::testing::HasSubstr("not found in message descriptor"));
+}
+
+}  // namespace
+
+}  // namespace tensorflow::tools::proto_splitter

tensorflow/tools/proto_splitter/python/BUILD

@@ -0,0 +1,68 @@
+# Description:
+#   Python library for splitting and joining large protos.
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "saved_model",
+    srcs = ["saved_model.py"],
+    # NOTE(yibaimeng): To be removed when everything is migrated to `pywrap_saved_model.Save`.
+    visibility = [
+        "//tensorflow:internal",
+        "//waymo/ml/deploy/tensorflow:__pkg__",
+    ],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/tools/proto_splitter:constants",
+        "//tensorflow/tools/proto_splitter:split",
+        "//tensorflow/tools/proto_splitter:split_graph_def",
+    ],
+)
+
+py_test(
+    name = "saved_model_test",
+    srcs = ["saved_model_test.py"],
+    strict_deps = True,
+    tags = [
+        "no_mac",  # b/291933687
+        "no_windows",  # b/291001524
+    ],
+    deps = [
+        ":saved_model",
+        ":test_util",
+        #internal proto upb dep
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/tools/proto_splitter:constants",
+    ],
+)
+
+pytype_strict_library(
+    name = "test_util",
+    srcs = ["test_util.py"],
+    deps = [
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_util",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_test(
+    name = "test_util_test",
+    srcs = ["test_util_test.py"],
+    strict_deps = True,
+    deps = [
+        #internal proto upb dep
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/tools/proto_splitter/python:test_util",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/tools/proto_splitter/python/saved_model.py

@@ -0,0 +1,47 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""SavedModel Splitter."""
+
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.tools.proto_splitter import constants
+from tensorflow.tools.proto_splitter import split
+from tensorflow.tools.proto_splitter import split_graph_def
+
+
+class SavedModelSplitter(split.ComposableSplitter):
+  """Splits a SavedModel proto into chunks of size < 2GB."""
+
+  def build_chunks(self):
+    if not isinstance(self._proto, saved_model_pb2.SavedModel):
+      raise TypeError(
+          "SavedModelSplitter can only split SavedModel protos. "
+          f"Got {type(self._proto)}."
+      )
+
+    if self._proto.ByteSize() >= constants.max_size():
+      graph_def = self._proto.meta_graphs[0].graph_def
+      graph_def_fields = ["meta_graphs", 0, "graph_def"]
+      split_graph_def.GraphDefSplitter(
+          self._proto.meta_graphs[0].graph_def,
+          parent_splitter=self,
+          fields_in_parent=graph_def_fields,
+      ).build_chunks()
+
+    # Check if the proto size is still larger than the max size.
+    if self._proto.ByteSize() >= constants.max_size():
+      # Create a chunk for the GraphDef, and ensure the GraphDef is merged in
+      # first by adding it at index 1. The 0th chunk is the SavedModel itself.
+      self.add_chunk(graph_def, graph_def_fields, index=1)
+      self._proto.meta_graphs[0].ClearField("graph_def")

tensorflow/tools/proto_splitter/python/saved_model_test.py

@@ -0,0 +1,56 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Tests for SavedModelSplitter."""
+
+
+from google.protobuf import message
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.python.platform import test
+from tensorflow.tools.proto_splitter import constants
+from tensorflow.tools.proto_splitter.python import saved_model
+from tensorflow.tools.proto_splitter.python import test_util
+
+
+class SavedModelSplitterTest(test.TestCase):
+
+  def _assert_chunk_sizes(self, chunks, max_size):
+    """Asserts that all chunk proto sizes are <= max_size."""
+    for chunk in chunks:
+      if isinstance(chunk, message.Message):
+        self.assertLessEqual(chunk.ByteSize(), max_size)
+
+  def test_split_saved_model(self):
+    sizes = [100, 100, 1000, 100, 1000, 500, 100, 100, 100]
+    fn1 = [100, 100, 100]
+    fn2 = [100, 500]
+    fn3 = [100]
+    fn4 = [100, 100]
+
+    max_size = 500
+    constants.debug_set_max_size(max_size)
+
+    graph_def = test_util.make_graph_def_with_constant_nodes(
+        sizes, fn1=fn1, fn2=fn2, fn3=fn3, fn4=fn4
+    )
+    proto = saved_model_pb2.SavedModel()
+    proto.meta_graphs.add().graph_def.CopyFrom(graph_def)
+
+    splitter = saved_model.SavedModelSplitter(proto)
+    chunks, _ = splitter.split()
+    self._assert_chunk_sizes(chunks, max_size)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/tools/proto_splitter/python/test_util.py

@@ -0,0 +1,78 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Utilities for Python tests."""
+
+from collections.abc import Sequence
+import math
+from typing import Optional
+
+import numpy as np
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_util
+
+
+def make_graph_def_with_constant_nodes(
+    node_sizes: Sequence[int],
+    dtype: Optional[dtypes.DType] = None,
+    **function_node_sizes,
+) -> graph_pb2.GraphDef:
+  """Creates a GraphDef with approximate node sizes.
+
+  Args:
+    node_sizes: list of ints, the approximate desired sizes of the nodes in the
+      GraphDef.
+    dtype: Dtype of encoded constant values (float32 or float64).
+    **function_node_sizes: Map of function name to FunctionDef node sizes (see
+      `node_sizes`).
+
+  Returns:
+    A GraphDef proto.
+  """
+  dtype = dtypes.float32
+  graph_def = graph_pb2.GraphDef()
+  n = 0
+
+  def add_nodes(node_list, sizes):
+    nonlocal n
+    for s in sizes:
+      node = node_list.add(name=f"Const_{n}", op="Const")
+
+      # Add an empty value to compute the approximate size of the constant
+      # value that will be added to the proto.
+      node.attr["value"].tensor.MergeFrom(
+          tensor_util.make_tensor_proto(np.ones([]), dtype=dtype)
+      )
+      remaining_size = s - node.ByteSize()
+      if remaining_size < 0:
+        raise ValueError(f"Unable to create node of size {s} bytes.")
+
+      constant_size = [math.ceil(remaining_size / dtype.size)]
+      node.attr["value"].tensor.Clear()
+      node.attr["value"].tensor.MergeFrom(
+          tensor_util.make_tensor_proto(
+              np.random.random_sample(constant_size), dtype=dtype
+          )
+      )
+      n += 1
+
+  add_nodes(graph_def.node, node_sizes)
+  for fn_name, fn_sizes in function_node_sizes.items():
+    fn = graph_def.library.function.add()
+    fn.signature.name = fn_name
+    add_nodes(fn.node_def, fn_sizes)
+
+  return graph_def

tensorflow/tools/proto_splitter/python/test_util_test.py

@@ -0,0 +1,51 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Tests for test_util."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.platform import test
+from tensorflow.tools.proto_splitter.python import test_util
+
+
+class MakeGraphDefTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ("Float64", dtypes.float64), ("Float32", dtypes.float32)
+  )
+  def testMakeGraphDef(self, dtype):
+    expected_sizes = [75, 50, 100, 95, 120]
+    fn1 = [121, 153, 250, 55]
+    fn2 = [552, 45]
+    graph_def = test_util.make_graph_def_with_constant_nodes(
+        expected_sizes, dtype=dtype, fn1=fn1, fn2=fn2)
+    self.assertAllClose(
+        expected_sizes, [node.ByteSize() for node in graph_def.node], atol=5
+    )
+    self.assertAllClose(
+        fn1,
+        [node.ByteSize() for node in graph_def.library.function[0].node_def],
+        atol=10,
+    )
+    self.assertAllClose(
+        fn2,
+        [node.ByteSize() for node in graph_def.library.function[1].node_def],
+        atol=10,
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/tools/proto_splitter/split.py

@@ -0,0 +1,291 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Basic interface for Python-based Splitter."""
+
+import abc
+from collections.abc import Sequence
+import time
+from typing import Optional, Union
+
+from absl import logging
+import riegeli
+
+from google.protobuf import message
+from tensorflow.python.lib.io import file_io
+from tensorflow.tools.proto_splitter import chunk_pb2
+from tensorflow.tools.proto_splitter import util
+from tensorflow.tools.proto_splitter import version as version_lib
+from tensorflow.tools.proto_splitter import versions_pb2
+
+
+class Splitter(abc.ABC):
+  """An abstract class for splitting and writing protos that are > 2GB.
+
+  See the README on how to use or subclass this class.
+  """
+
+  @property
+  @abc.abstractmethod
+  def version_def(self) -> versions_pb2.VersionDef:
+    """Version info about the splitter and merge implementation required."""
+
+  @abc.abstractmethod
+  def split(
+      self,
+  ) -> tuple[Sequence[Union[message.Message, bytes]], chunk_pb2.ChunkedMessage]:
+    """Splits proto message into a Sequence of protos/bytes."""
+
+  @abc.abstractmethod
+  def write(self, file_prefix: str) -> str:
+    """Serializes proto to disk.
+
+    Args:
+      file_prefix: string prefix of the filepath.
+
+    Returns:
+      The actual path the proto is written to.
+    """
+
+
+class ComposableSplitter(Splitter):
+  """A Splitter that can be composed with other splitters.
+
+  This Splitter writes to the riegeli file format.
+
+  See README for details.
+  """
+
+  def __init__(
+      self,
+      proto,
+      *,
+      proto_as_initial_chunk: bool = True,
+      parent_splitter: Optional["ComposableSplitter"] = None,
+      fields_in_parent: Optional[util.FieldTypes] = None,
+  ):
+    """Initializes ComposableSplitter.
+
+    Args:
+      proto: Proto message to split.
+      proto_as_initial_chunk: Whether to initialize chunks with the
+        user-provided proto as the initial chunk.
+      parent_splitter: The parent `ComposableSplitter` object.
+      fields_in_parent: Fields to access `proto` from the parent splitter's
+        proto.
+    """
+    self._proto = proto
+    self._parent_splitter = parent_splitter
+    self._fields_in_parent = fields_in_parent
+
+    # Whether chunks have been created. See `build_chunks()`.
+    self._built = False
+
+    # Keep a list of chunk ids in the order in which they were added to the
+    # list.
+    self._add_chunk_order = []
+    self._fix_chunk_order = False
+
+    # Initialize chunks and ChunkedMessage (optionally with the first chunk as
+    # the user-provided proto.
+    if parent_splitter is not None:
+      # If this is not the root Splitter class, skip the initialization of
+      # the chunks/message since the parent's will be updated instead.
+      self._chunks = None
+      self._chunked_message = None
+    elif proto_as_initial_chunk:
+      self._chunks = [self._proto]
+      self._chunked_message = chunk_pb2.ChunkedMessage(chunk_index=0)
+      self._add_chunk_order.append(id(self._proto))
+    else:
+      self._chunks = []
+      self._chunked_message = chunk_pb2.ChunkedMessage()
+
+  def build_chunks(self) -> None:
+    """Builds the Splitter object by generating chunks from the proto.
+
+    Subclasses of `ComposableChunks` should only need to override this method.
+
+    This method should be called once per Splitter to create the chunks.
+    Users should call the methods `split` or `write` instead.
+    """
+
+  @property
+  def version_def(self) -> versions_pb2.VersionDef:
+    """Version info about the splitter and join implementation required."""
+    return versions_pb2.VersionDef(
+        splitter_version=1,
+        join_version=0,
+        bad_consumers=version_lib.get_bad_versions(),
+    )
+
+  def split(
+      self,
+  ) -> tuple[Sequence[Union[message.Message, bytes]], chunk_pb2.ChunkedMessage]:
+    """Splits a proto message into a Sequence of protos/bytes."""
+    if self._parent_splitter:
+      raise ValueError(
+          "A child ComposableSplitter's `split` method should not be called "
+          "directly, since it inherit chunks from a parent object. Please call "
+          "the parent's `split()` method instead."
+      )
+
+    assert self._chunks is not None
+    assert self._chunked_message is not None
+
+    if not self._built:
+      self.build_chunks()
+      self._fix_chunks()
+      self._built = True
+    return self._chunks, self._chunked_message
+
+  def write(
+      self, file_prefix: str, writer_options: Optional[str] = None
+  ) -> str:
+    """Serializes a proto to disk.
+
+    The writer writes all chunks into a riegeli file. The chunk metadata
+    (ChunkMetadata) is written at the very end.
+
+    Args:
+      file_prefix: string prefix of the filepath. The writer will automatically
+        attach a `.pb` or `.cpb` (chunked pb) suffix depending on whether the
+        proto is split.
+      writer_options: Optional writer options to pass to the riegeli writer. See
+        https://github.com/google/riegeli/blob/master/doc/record_writer_options.md
+        for options.
+
+    Returns:
+      The actual filepath the proto is written to. The filepath will be
+      different depending on whether the proto is split, i.e., whether it will
+      be a pb or not.
+    """
+    if self._parent_splitter is not None:
+      raise ValueError(
+          "A child ComposableSplitter's `write` method should not be called "
+          "directly, since it inherits unrelated chunks from a parent object. "
+          "Please call the parent's `write()` method instead."
+      )
+
+    start_time = time.time()
+    chunks, chunked_message = self.split()
+
+    if not chunked_message.chunked_fields:
+      path = f"{file_prefix}.pb"
+      file_io.atomic_write_string_to_file(
+          path, self._proto.SerializeToString(deterministic=True)
+      )
+      logging.info("Unchunked file exported to %s", path)
+      return path
+
+    path = f"{file_prefix}.cpb"
+    writer_kwargs = {}
+    if writer_options is not None:
+      writer_kwargs["options"] = writer_options
+    with riegeli.RecordWriter(file_io.FileIO(path, "wb"), **writer_kwargs) as f:
+      metadata = chunk_pb2.ChunkMetadata(
+          message=chunked_message, version=self.version_def
+      )
+      for chunk in chunks:
+        if isinstance(chunk, message.Message):
+          f.write_message(chunk)
+          chunk_type = chunk_pb2.ChunkInfo.Type.MESSAGE
+          size = chunk.ByteSize()
+        else:
+          f.write_record(chunk)
+          chunk_type = chunk_pb2.ChunkInfo.Type.BYTES
+          size = len(chunk)
+        metadata.chunks.add(
+            type=chunk_type, size=size, offset=f.last_pos.numeric
+        )
+      f.write_message(metadata)
+
+    end = time.time()
+
+    logging.info("Chunked file exported to %s", path)
+    logging.info(
+        "Total time spent splitting and writing the message: %s",
+        end - start_time,
+    )
+    logging.info(
+        "Number of chunks created (including initial message): %s",
+        len(chunks),
+    )
+    return path
+
+  def add_chunk(
+      self,
+      chunk: Union[message.Message, bytes],
+      field_tags: util.FieldTypes,
+      index=None,
+  ) -> None:
+    """Adds a new chunk and updates the ChunkedMessage proto.
+
+    Args:
+      chunk: Proto message or bytes.
+      field_tags: Field information about the placement of the chunked data
+        within self._proto.
+      index: Optional index at which to insert the chunk. The chunk ordering is
+        important for merging.
+    """
+    if self._parent_splitter is not None:
+      self._parent_splitter.add_chunk(
+          chunk, self._fields_in_parent + field_tags, index
+      )
+    else:
+      assert self._chunks is not None
+      assert self._chunked_message is not None
+      field = self._chunked_message.chunked_fields.add(
+          field_tag=util.get_field_tag(self._proto, field_tags)
+      )
+      new_chunk_index = len(self._chunks)
+      field.message.chunk_index = new_chunk_index
+      self._add_chunk_order.append(id(chunk))
+
+      if index is None:
+        self._chunks.append(chunk)
+      else:
+        self._chunks.insert(index, chunk)
+        self._fix_chunk_order = True
+
+  def _fix_chunks(self) -> None:
+    """Fixes chunk indices in the ChunkedMessage."""
+    if not self._fix_chunk_order:
+      return
+
+    # The chunk_index of each nested ChunkedMessage is set to the length of the
+    # list when the chunk was added. This would be fine if the chunks were
+    # always added to the end of the list. However, this is not always the case
+    # the indices must be updated.
+
+    # Use the address of each chunk (python `id`) as lookup keys to the
+    # ordered chunk indices.
+    chunk_indices = {id(chunk): i for i, chunk in enumerate(self._chunks)}
+
+    to_fix = [self._chunked_message]
+    while to_fix:
+      for field in to_fix.pop().chunked_fields:
+        if field.message.chunked_fields:
+          to_fix.append(field.message)
+        if not field.message.HasField("chunk_index"):
+          continue
+        chunk_addr = self._add_chunk_order[field.message.chunk_index]
+        assert (
+            chunk_addr in chunk_indices
+        ), f"Found unexpected chunk {chunk_addr}"
+        new_chunk_index = chunk_indices[chunk_addr]
+        field.message.chunk_index = new_chunk_index
+
+    self._add_chunk_order = [id(chunk) for chunk in self._chunks]
+    self._fix_chunk_order = False

tensorflow/tools/proto_splitter/split_graph_def.py

@@ -0,0 +1,329 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""GraphDef splitter."""
+
+from collections.abc import Sequence
+import itertools
+from typing import Optional, Type
+
+from google.protobuf import message
+from tensorflow.core.framework import graph_pb2
+from tensorflow.core.framework import node_def_pb2
+from tensorflow.python.framework import tensor_util
+from tensorflow.tools.proto_splitter import constants
+from tensorflow.tools.proto_splitter import split
+from tensorflow.tools.proto_splitter import util
+
+_CONST_OP = "Const"
+
+
+class GraphDefSplitter(split.ComposableSplitter):
+  """Implements proto splitter for GraphDef.
+
+  This Splitter will modify the passed in proto in place.
+  """
+
+  def build_chunks(self):
+    """Splits a GraphDef proto into smaller chunks."""
+    proto = self._proto
+    if not isinstance(proto, graph_pb2.GraphDef):
+      raise TypeError("Can only split GraphDef type protos.")
+
+    proto_size = proto.ByteSize()
+    if proto_size < constants.max_size():
+      return
+
+    # Split `GraphDef.node`
+    node_splitter = RepeatedMessageSplitter(
+        proto,
+        "node",
+        [ConstantNodeDefSplitter, LargeMessageSplitter],
+        parent_splitter=self,
+        fields_in_parent=[],
+    )
+
+    # Split `GraphDef.library.function`
+    function_splitter = RepeatedMessageSplitter(
+        proto.library,
+        ["function"],
+        [FunctionDefSplitter],
+        parent_splitter=self,
+        fields_in_parent=["library"],
+    )
+
+    library_size = proto.library.ByteSize()
+    approx_node_size = proto_size - library_size
+
+    if library_size > approx_node_size:
+      library_size -= function_splitter.build_chunks()
+      if library_size + approx_node_size > constants.max_size():
+        approx_node_size -= node_splitter.build_chunks()
+
+    else:
+      approx_node_size -= node_splitter.build_chunks()
+      if library_size + approx_node_size > constants.max_size():
+        library_size -= function_splitter.build_chunks()
+
+    if proto.ByteSize() > constants.max_size():
+      # Since there are chunks with the "library" field tag, insert this
+      # chunk before the other chunks at index 1 (index 0 is reserved for the
+      # base chunk).
+      self.add_chunk(proto.library, ["library"], 1)
+      proto.ClearField("library")
+
+_KEEP_TENSOR_PROTO_FIELDS = ("dtype", "tensor_shape", "version_number")
+
+
+def chunk_constant_value(node: node_def_pb2.NodeDef, size: int):
+  """Extracts and clears the constant value from a NodeDef.
+
+  Args:
+    node: NodeDef with const value to extract.
+    size: Size of NodeDef (for error reporting).
+
+  Returns:
+    Bytes representation of the Constant tensor content.
+  """
+  if node.op == _CONST_OP:
+    tensor_proto = node.attr["value"].tensor
+    if tensor_proto.tensor_content:
+      b = tensor_proto.tensor_content
+    else:
+      # The raw tensor value could be stored in one of the "xxx_val" attributes.
+      # Extract it here, and convert to bytes.
+      b = tensor_util.MakeNdarray(tensor_proto).tobytes()
+
+    # Keep the TensorProto's dtype, tensor_shape, and version_number fields,
+    # but clear the raw tensor content / "xxx_val" attributes.
+    kept_attributes = {
+        key: getattr(tensor_proto, key) for key in _KEEP_TENSOR_PROTO_FIELDS
+    }
+    tensor_proto.Clear()
+    for field, val in kept_attributes.items():
+      if isinstance(val, message.Message):
+        getattr(tensor_proto, field).MergeFrom(val)
+      else:
+        setattr(tensor_proto, field, val)
+
+    return b
+  else:
+    attributes_and_sizes = ", ".join(
+        [
+            f"{key}: {util.format_bytes(val.ByteSize())}"
+            for key, val in node.attr.items()
+        ]
+    )
+    raise ValueError(
+        "Unable to split GraphDef because at least one of the nodes "
+        "individually exceeds the max size of "
+        f"{util.format_bytes(constants.max_size())}. "
+        "Currently only Const nodes can be further split."
+        "\nNode info:"
+        f"\n\tsize: {util.format_bytes(size)}"
+        f"\n\tname: {node.name}"
+        f"\n\top: {node.op}"
+        f"\n\tinputs: {node.input}"
+        f"\n\top: {node.op}"
+        f"\n\tdevice: {node.device}"
+        f"\n\tattr (and sizes): {attributes_and_sizes}"
+    )
+
+
+def _split_repeated_field(
+    proto: message.Message,
+    new_proto: message.Message,
+    fields: util.FieldTypes,
+    start_index: int,
+    end_index: Optional[int] = None,
+) -> None:
+  """Generic function for copying a repeated field from one proto to another."""
+  util.get_field(new_proto, fields)[0].MergeFrom(
+      util.get_field(proto, fields)[0][start_index:end_index]
+  )
+
+
+_ABOVE_MAX_SIZE = lambda x: x > constants.max_size()
+_GREEDY_SPLIT = lambda x: x > constants.max_size() // 3
+_ALWAYS_SPLIT = lambda x: True
+
+
+class SplitBasedOnSize(split.ComposableSplitter):
+  """A Splitter that's based on the size of the input proto."""
+
+  __slots__ = ("fn", "proto_size")
+
+  def __init__(self, proto, proto_size, **kwargs):
+    """Initializer."""
+    self.proto_size = proto_size
+    super().__init__(proto, **kwargs)
+
+  def build_chunks(self) -> int:
+    """Splits the proto, and returns the size of the chunks created."""
+    return 0
+
+
+class RepeatedMessageSplitter(split.ComposableSplitter):
+  """Splits a repeated message field on a proto."""
+
+  __slots__ = ("repeated_field", "message_splitters")
+
+  def __init__(
+      self,
+      proto,
+      repeated_field: util.FieldTypes,
+      message_splitters: Sequence[Type[SplitBasedOnSize]],
+      **kwargs,
+  ):
+    """Initializer."""
+    super().__init__(proto, **kwargs)
+    if not isinstance(repeated_field, list):
+      repeated_field = [repeated_field]
+    self.repeated_field = repeated_field
+    self.message_splitters = message_splitters
+
+  def build_chunks(self) -> int:
+    """Splits the proto, and returns the size of the chunks created."""
+    proto = self._proto
+
+    total_size_diff = 0
+
+    field, field_desc = util.get_field(proto, self.repeated_field)
+    if not util.is_repeated(field_desc) and field_desc.message_type:
+      raise ValueError(
+          "RepeatedMessageSplitter can only be used on repeated fields. "
+          f"Got proto={type(proto)}, field='{field_desc.name}'"
+      )
+
+    # List of indices at which to split the repeated field. For example, [3, 5]
+    # means that the field list is split into: [:3], [3:5], [5:]
+    repeated_msg_split = []
+    # Should be the same length as the list above. Contains new protos to hold
+    # the elements that are split from the original proto.
+    # From the [3, 5] example above, the messages in this list contain nodes
+    # [3:5] and [5:]
+    repeated_msg_graphs = []
+    # Track the total size of the current node split.
+    total_size = 0
+
+    # Linearly iterate through all nodes. It may be possible to optimize this
+    # further by making best guesses as to where to split the nodes, since
+    # most nodes (aside from constants) are relatively small.
+    for n, ele in enumerate(field):
+      size = ele.ByteSize()
+
+      for splitter_cls in self.message_splitters:
+        splitter = splitter_cls(
+            ele,
+            size,
+            parent_splitter=self,
+            fields_in_parent=self.repeated_field + [n],
+        )
+        size_diff = splitter.build_chunks()
+        total_size_diff += size_diff
+        size -= size_diff
+
+      # Create a new GraphDef chunk if the current list of nodes is too large.
+      if total_size + size >= constants.max_size():
+        new_msg = type(self._proto)()
+        repeated_msg_split.append(n)
+        repeated_msg_graphs.append(new_msg)
+        self.add_chunk(new_msg, [])
+
+        if len(repeated_msg_split) >= 1:
+          total_size_diff += total_size
+
+        total_size = 0
+
+      total_size += size
+
+    if repeated_msg_split:
+      # Finish writing repeated chunks.
+      start = repeated_msg_split[0]
+      for end, msg in zip(
+          itertools.chain.from_iterable([repeated_msg_split[1:], [None]]),
+          repeated_msg_graphs,
+      ):
+        _split_repeated_field(proto, msg, self.repeated_field, start, end)
+        start = end
+      del field[repeated_msg_split[0] :]
+
+    return total_size_diff
+
+
+class ConstantNodeDefSplitter(SplitBasedOnSize):
+  """Extracts constant value from a `Const` NodeDef."""
+
+  def build_chunks(self) -> int:
+    """Splits a NodeDef proto, and returns the size of the chunks created."""
+    if _ABOVE_MAX_SIZE(self.proto_size):
+      constant_bytes = chunk_constant_value(self._proto, self.proto_size)
+      self.add_chunk(
+          constant_bytes,
+          ["attr", "value", "tensor", "tensor_content"],
+      )
+      return len(constant_bytes)
+    return 0
+
+
+class LargeMessageSplitter(SplitBasedOnSize):
+  """Splits a message into a separaet chunk if its over a certain size."""
+
+  __slots__ = ("size_check",)
+
+  def __init__(self, proto, proto_size, size_check=_GREEDY_SPLIT, **kwargs):
+    """Initializer."""
+
+    self.size_check = size_check
+    super().__init__(proto, proto_size, **kwargs)
+
+  def build_chunks(self) -> int:
+    """Creates a chunk for the entire proto and returns the original size."""
+    if self.size_check(self.proto_size):
+      new_proto = type(self._proto)()
+      new_proto.MergeFrom(self._proto)
+      self._proto.Clear()
+      self.add_chunk(new_proto, [])
+      return self.proto_size
+    return 0
+
+
+class FunctionDefSplitter(SplitBasedOnSize):
+  """Splits the FunctionDef message type."""
+
+  def build_chunks(self) -> int:
+    """Splits the proto, and returns the size of the chunks created."""
+    size_diff = 0
+
+    # First check if the entire FunctionDef can be split into a separate chunk.
+    # We do this before the `RepeatedMessageSplitter`, which is costly because
+    # it iterates through every `node_def`.
+    if _GREEDY_SPLIT(self.proto_size) and not _ABOVE_MAX_SIZE(self.proto_size):
+      size_diff += LargeMessageSplitter(
+          self._proto,
+          self.proto_size,
+          parent_splitter=self,
+          fields_in_parent=[],
+      ).build_chunks()
+
+    if _ABOVE_MAX_SIZE(self.proto_size):
+      # Split FunctionDefLibrary.function.node_def
+      size_diff += RepeatedMessageSplitter(
+          self._proto,
+          "node_def",
+          [ConstantNodeDefSplitter, LargeMessageSplitter],
+          parent_splitter=self,
+          fields_in_parent=[],
+      ).build_chunks()
+    return size_diff

tensorflow/tools/proto_splitter/split_graph_def_test.py

@@ -0,0 +1,294 @@
+# Copyright 2023 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.
+#
+# ==============================================================================
+"""Tests for GraphDef splitter."""
+
+import itertools
+
+from google.protobuf import message
+from google.protobuf import text_format
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.framework import graph_pb2
+from tensorflow.python.platform import test
+from tensorflow.tools.proto_splitter import chunk_pb2
+from tensorflow.tools.proto_splitter import constants
+from tensorflow.tools.proto_splitter import split_graph_def
+from tensorflow.tools.proto_splitter import util
+from tensorflow.tools.proto_splitter.python import test_util
+
+
+class GraphDefSplitterTest(test.TestCase):
+
+  def _make_graph_def_with_constant_nodes(
+      self, node_sizes, dtype=None, **function_node_sizes
+  ):
+    return test_util.make_graph_def_with_constant_nodes(
+        node_sizes, dtype, **function_node_sizes
+    )
+
+  def _copy_graph(self, graph_def):
+    """Create a copy of GraphDef."""
+    graph_def_copy = graph_pb2.GraphDef()
+    graph_def_copy.CopyFrom(graph_def)
+    return graph_def_copy
+
+  def _assert_chunk_sizes(self, chunks, max_size):
+    """Asserts that all chunk proto sizes are <= max_size."""
+    for chunk in chunks:
+      if isinstance(chunk, message.Message):
+        self.assertLessEqual(chunk.ByteSize(), max_size)
+
+  def _assert_field_tags(self, expected_fields, actual_fields):
+    self.assertLen(actual_fields, len(expected_fields))
+    for expected, actual in zip(expected_fields, actual_fields):
+      self.assertProtoEquals(expected, actual)
+
+  def testSplitNoChunks(self):
+    sizes = [50, 100, 50, 50, 100]
+    max_size = 500
+    constants.debug_set_max_size(max_size)
+
+    graph_def = self._make_graph_def_with_constant_nodes(sizes)
+    s = split_graph_def.GraphDefSplitter(self._copy_graph(graph_def))
+    chunks, _ = s.split()
+    self.assertLen(chunks, 1)
+    self.assertProtoEquals(graph_def, chunks[0])
+
+  def testLargeConstant(self):
+    sizes = [50, 50, 1000, 50, 1000]
+    max_size = 500
+    constants.debug_set_max_size(max_size)
+
+    #    Expected Chunks (Max Size = 500)
+    #    -----------------------------
+    #       Chunk #: Contents
+    #    -----------------------------
+    #       0: GraphDef
+    #    -----------------------------
+    #       1: GraphDef.nodes[2].attr["value"].tensor.tensor_content
+    #    -----------------------------
+    #       2: GraphDef.nodes[4].attr["value"].tensor.tensor_content
+    #    -----------------------------
+
+    graph_def = self._make_graph_def_with_constant_nodes(sizes)
+    s = split_graph_def.GraphDefSplitter(self._copy_graph(graph_def))
+    chunks, chunked_message = s.split()
+    self.assertLen(chunks, 3)
+    self._assert_chunk_sizes(chunks, max_size)
+
+    self.assertEqual(
+        graph_def.node[4].attr["value"].tensor.tensor_content, chunks[2]
+    )
+    self.assertEqual(
+        graph_def.node[2].attr["value"].tensor.tensor_content, chunks[1]
+    )
+
+    # Check the ChunkedMessage proto.
+    self.assertLen(chunked_message.chunked_fields, 2)
+    self.assertEqual(1, chunked_message.chunked_fields[0].message.chunk_index)
+    self.assertEqual(2, chunked_message.chunked_fields[1].message.chunk_index)
+    self._assert_field_tags(
+        util.get_field_tag(
+            graph_def, ["node", 2, "attr", "value", "tensor", "tensor_content"]
+        ),
+        chunked_message.chunked_fields[0].field_tag,
+    )
+    self._assert_field_tags(
+        util.get_field_tag(
+            graph_def, ["node", 4, "attr", "value", "tensor", "tensor_content"]
+        ),
+        chunked_message.chunked_fields[1].field_tag,
+    )
+
+  def testLotsOfNodes(self):
+    # The actual sizes in the generated graph has a slight deviation, but are
+    # between [90, 100] (tested in testMakeGraphDef with atol=5).
+    sizes = [95] * 15
+    max_size = 500
+    constants.debug_set_max_size(500)
+
+    #    Expected Chunks (Max Size = 500)
+    #    -----------------------------
+    #       Chunk #: Contents
+    #    -----------------------------
+    #       0: GraphDef  # (nodes [0:5])
+    #    -----------------------------
+    #       1: GraphDef  # (nodes [5:10])
+    #    -----------------------------
+    #       2: GraphDef  # (nodes [10:15])
+    #    -----------------------------
+    graph_def = self._make_graph_def_with_constant_nodes(sizes)
+    s = split_graph_def.GraphDefSplitter(self._copy_graph(graph_def))
+    chunks, chunked_message = s.split()
+    self.assertLen(chunks, 3)
+    self._assert_chunk_sizes(chunks, max_size)
+    for node, chunk in zip(
+        graph_def.node,
+        itertools.chain(chunks[0].node, chunks[1].node, chunks[2].node),
+    ):
+      self.assertProtoEquals(node, chunk)
+
+    # Check the ChunkedMessage proto.
+    self.assertLen(chunked_message.chunked_fields, 2)
+    self.assertEqual(1, chunked_message.chunked_fields[0].message.chunk_index)
+    self.assertEqual(2, chunked_message.chunked_fields[1].message.chunk_index)
+    self.assertEmpty(chunked_message.chunked_fields[0].field_tag)
+    self.assertEmpty(chunked_message.chunked_fields[1].field_tag)
+
+  def testLargeNodes(self):
+    # Large nodes are greedily split from the original proto if they are
+    # larger than max_size / 3.
+    sizes = [50, 95, 95, 95, 50, 95]
+    max_size = 200
+    constants.debug_set_max_size(max_size)
+
+    # This should create 6 chunks:
+    #   [parent GraphDef, node[1], node[2], node[3], node[5], ChunkedMessage]
+    graph_def = self._make_graph_def_with_constant_nodes(sizes)
+    s = split_graph_def.GraphDefSplitter(self._copy_graph(graph_def))
+    chunks, chunked_message = s.split()
+    self.assertLen(chunks, 5)
+    self._assert_chunk_sizes(chunks, max_size)
+
+    # Check large nodes are chunked away.
+    self.assertProtoEquals(graph_def.node[1], chunks[1])
+    self.assertProtoEquals(graph_def.node[2], chunks[2])
+    self.assertProtoEquals(graph_def.node[3], chunks[3])
+    self.assertProtoEquals(graph_def.node[5], chunks[4])
+
+    # Check that the parent GraphDef still contains small nodes.
+    self.assertProtoEquals(graph_def.node[0], chunks[0].node[0])
+    self.assertProtoEquals(graph_def.node[4], chunks[0].node[4])
+
+    # Check that the parent GraphDef contains empty nodes where the large nodes
+    # were originally.
+    self.assertEqual(0, chunks[0].node[1].ByteSize())
+    self.assertEqual(0, chunks[0].node[2].ByteSize())
+    self.assertEqual(0, chunks[0].node[3].ByteSize())
+    self.assertEqual(0, chunks[0].node[5].ByteSize())
+
+    # Check the ChunkedMessage proto.
+    self.assertLen(chunked_message.chunked_fields, 4)
+    self._assert_field_tags(
+        util.get_field_tag(graph_def, ["node", 1]),
+        chunked_message.chunked_fields[0].field_tag,
+    )
+    self._assert_field_tags(
+        util.get_field_tag(graph_def, ["node", 2]),
+        chunked_message.chunked_fields[1].field_tag,
+    )
+    self._assert_field_tags(
+        util.get_field_tag(graph_def, ["node", 3]),
+        chunked_message.chunked_fields[2].field_tag,
+    )
+    self._assert_field_tags(
+        util.get_field_tag(graph_def, ["node", 5]),
+        chunked_message.chunked_fields[3].field_tag,
+    )
+
+  def testFunctionLotsOfNodes(self):
+    sizes = []
+    fn1 = [50, 50, 50, 50, 50]
+    max_size = 200
+    constants.debug_set_max_size(max_size)
+
+    graph_def = self._make_graph_def_with_constant_nodes(sizes, fn=fn1)
+    s = split_graph_def.GraphDefSplitter(self._copy_graph(graph_def))
+    chunks, chunked_message = s.split()
+
+    self.assertLen(chunks, 2)
+    self.assertIsInstance(chunks[0], graph_pb2.GraphDef)
+    self.assertIsInstance(chunks[1], function_pb2.FunctionDef)
+    self._assert_chunk_sizes(chunks, max_size)
+
+    for node, chunk in zip(
+        graph_def.library.function[0].node_def,
+        itertools.chain(
+            chunks[0].library.function[0].node_def, chunks[1].node_def
+        ),
+    ):
+      self.assertProtoEquals(node, chunk)
+
+    expected_message = chunk_pb2.ChunkedMessage()
+    text_format.Parse(
+        """
+        chunk_index: 0
+        chunked_fields {
+            field_tag {
+                field: 2
+            }
+            field_tag {
+                field: 1
+            }
+            field_tag {
+                index: 0
+            }
+            message {
+                chunk_index: 1
+            }
+        }""",
+        expected_message,
+    )
+    self.assertProtoEquals(expected_message, chunked_message)
+
+  def testFunctionLargeNodes(self):
+    sizes = []
+    fn1 = [500, 500, 50, 500]
+    max_size = 200
+    constants.debug_set_max_size(max_size)
+
+    graph_def = self._make_graph_def_with_constant_nodes(sizes, fn=fn1)
+    s = split_graph_def.GraphDefSplitter(self._copy_graph(graph_def))
+    chunks, _ = s.split()
+
+    self.assertLen(chunks, 4)
+    self._assert_chunk_sizes(chunks, max_size)
+    self.assertIsInstance(chunks[0], graph_pb2.GraphDef)
+
+    def get_const_value(index):
+      node_def = graph_def.library.function[0].node_def[index]
+      return node_def.attr["value"].tensor.tensor_content
+
+    expected_values = [
+        get_const_value(0),
+        get_const_value(1),
+        get_const_value(3),
+    ]
+    for expected, chunk in zip(expected_values, chunks[1:]):
+      self.assertEqual(expected, chunk)
+
+  def testChunkGraphDefAndFunctions(self):
+    sizes = [50, 50, 50, 50, 50, 50]
+    fn1 = [50, 50, 50]
+    fn2 = [50]
+    fn3 = [50]
+    fn4 = [50]
+    max_size = 200
+    constants.debug_set_max_size(max_size)
+
+    graph_def = self._make_graph_def_with_constant_nodes(
+        sizes, fn1=fn1, fn2=fn2, fn3=fn3, fn4=fn4
+    )
+    s = split_graph_def.GraphDefSplitter(self._copy_graph(graph_def))
+    chunks, _ = s.split()
+
+    # Expected chunks:
+    # GraphDef.nodes[:3], GraphDef.nodes[3:], fn1, FunctionDefLibrary
+    self.assertLen(chunks, 4)
+    self._assert_chunk_sizes(chunks, max_size)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/tools/proto_splitter/split_test.py

@@ -0,0 +1,168 @@
+# Copyright 2023 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.
+#
+# ==============================================================================
+"""Tests for ProtoSplitter."""
+
+import os
+import random
+import string
+from typing import Optional
+
+from absl.testing import parameterized
+import riegeli
+
+from tensorflow.python.platform import test
+from tensorflow.tools.proto_splitter import chunk_pb2
+from tensorflow.tools.proto_splitter import split
+from tensorflow.tools.proto_splitter.testdata import test_message_pb2
+
+
+class RepeatedStringSplitter(split.ComposableSplitter):
+  """Splits a RepeatedString proto with N repeated strings into N chunks."""
+
+  def __init__(self, proto, **kwargs):
+    if not isinstance(proto, test_message_pb2.RepeatedString):
+      raise TypeError("Can only split RepeatedString type protos")
+
+    empty_strings = not bool(proto.strings)
+    super().__init__(proto, proto_as_initial_chunk=empty_strings, **kwargs)
+
+  def build_chunks(self):
+    for n, s in enumerate(self._proto.strings):
+      b = bytes(s, encoding="utf-8")
+      self.add_chunk(b, ["strings", n])
+    self._proto.ClearField("strings")
+
+
+def _random_string(length):
+  return bytes(
+      "".join(random.choices(string.ascii_lowercase, k=length)),
+      encoding="utf-8",
+  )
+
+
+class SplitRepeatedStringTest(test.TestCase):
+
+  def _to_proto(self, strings):
+    return test_message_pb2.RepeatedString(strings=strings)
+
+  def testSplit(self):
+    s = RepeatedStringSplitter(test_message_pb2.RepeatedString(strings=[]))
+    chunks = s.split()[0]
+    self.assertLen(chunks, 1)
+    self.assertIsInstance(chunks[0], test_message_pb2.RepeatedString)
+
+    s = RepeatedStringSplitter(
+        test_message_pb2.RepeatedString(strings=["a", "b", "c"])
+    )
+    chunks, chunked_message = s.split()
+    self.assertListEqual([b"a", b"b", b"c"], chunks)
+    self.assertLen(chunked_message.chunked_fields, 3)
+
+  def testWrite(self):
+    path = os.path.join(self.create_tempdir(), "split-repeat")
+    data = [_random_string(5), _random_string(10), _random_string(15)]
+    returned_path = RepeatedStringSplitter(
+        test_message_pb2.RepeatedString(strings=data)
+    ).write(path)
+    self.assertEqual(returned_path, f"{path}.cpb")
+
+    with riegeli.RecordReader(open(f"{path}.cpb", "rb")) as reader:
+      self.assertTrue(reader.check_file_format())
+      records = list(reader.read_records())
+      self.assertLen(records, 4)
+
+      proto = chunk_pb2.ChunkMetadata()
+      proto.ParseFromString(records[-1])
+      self.assertLen(proto.message.chunked_fields, 3)
+      self.assertFalse(proto.message.HasField("chunk_index"))
+
+      expected_indices = [0, 1, 2]
+      # Check that the chunk indices and info are correct.
+      for expected_index, expected_data, chunk in zip(
+          expected_indices, data, proto.message.chunked_fields
+      ):
+        i = chunk.message.chunk_index
+        self.assertEqual(expected_index, i)
+
+        chunk_info = proto.chunks[i]
+        self.assertEqual(chunk_pb2.ChunkInfo.Type.BYTES, chunk_info.type)
+        self.assertEqual(len(expected_data), chunk_info.size)
+        reader.seek_numeric(chunk_info.offset)
+        self.assertEqual(expected_data, reader.read_record())
+
+  def test_child_splitter(self):
+    proto = test_message_pb2.RepeatedRepeatedString(
+        rs=[
+            test_message_pb2.RepeatedString(strings=["a", "b", "c"]),
+            test_message_pb2.RepeatedString(strings=["d", "e"]),
+        ]
+    )
+    splitter = NoOpSplitter(proto)
+
+    self.assertLen(splitter.split()[0], 1)
+    splitter.add_chunk("", [])
+    self.assertLen(splitter.split()[0], 2)
+
+    child = RepeatedStringSplitter(
+        proto.rs[0], parent_splitter=splitter, fields_in_parent=["rs", 0]
+    )
+    child.build_chunks()
+
+    self.assertLen(splitter.split()[0], 5)  # Adds 3 chunks.
+
+    RepeatedStringSplitter(
+        proto.rs[1], parent_splitter=splitter, fields_in_parent=["rs", 1]
+    ).build_chunks()
+
+    self.assertLen(splitter.split()[0], 7)  # Adds 2 chunks.
+
+    with self.assertRaisesRegex(
+        ValueError, " `split` method should not be called directly"
+    ):
+      child.split()
+    path = os.path.join(self.create_tempdir(), "na-split")
+    with self.assertRaisesRegex(
+        ValueError, " `write` method should not be called directly"
+    ):
+      child.write(path)
+
+
+class NoOpSplitter(split.ComposableSplitter):
+
+  def build_chunks(self):
+    pass
+
+
+class NoOpSplitterTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters([None, "uncompressed"])
+  def testWriteNoChunks(self, writer_options: Optional[str]):
+    path = os.path.join(self.create_tempdir(), "split-none")
+    proto = test_message_pb2.RepeatedString(strings=["a", "bc", "de"])
+    returned_path = NoOpSplitter(proto).write(path, writer_options)
+
+    expected_file_path = path + ".pb"
+    self.assertTrue(os.path.isfile(expected_file_path))
+    self.assertEqual(returned_path, expected_file_path)
+
+    parsed_proto = test_message_pb2.RepeatedString()
+    with open(expected_file_path, "rb") as f:
+      parsed_proto.ParseFromString(f.read())
+    self.assertProtoEquals(proto, parsed_proto)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/tools/proto_splitter/testdata/BUILD

@@ -0,0 +1,121 @@
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/tools/proto_splitter:__subpackages__"],
+    licenses = ["notice"],
+)
+
+exports_files([
+    "df-split-tree.cpb",
+    "df-split-tree.pbtxt",
+    "bf-split-tree.cpb",
+    "bf-split-tree.pbtxt",
+    "split-tree.pbtxt",
+    "split-large-nodes.cpb",
+    "split-large-nodes.pb",
+    "split-large-nodes.pbtxt",
+    "split-standard.cpb",
+    "split-standard.pb",
+    "split-standard.pbtxt",
+    "split-large-constant.cpb",
+    "split-large-constant.pbtxt",
+    "split-lots-nodes.cpb",
+    "split-lots-nodes.pb",
+    "split-lots-nodes.pbtxt",
+    "split-large-constant.cpb",
+    "split-large-constant.pb",
+    "split-large-constant.pbtxt",
+    "function-large-nodes.pb",
+    "function-lots-of-nodes.pb",
+    "graph-def-and-function.pb",
+    "many-field.cpb",
+    "many-field.pbtxt",
+])
+
+tf_proto_library(
+    name = "test_message_proto",
+    srcs = ["test_message.proto"],
+    make_default_target_header_only = True,
+    visibility = [
+        "//tensorflow/cc/saved_model:__subpackages__",
+        "//tensorflow/tools/proto_splitter:__subpackages__",
+    ],
+)
+
+# copybara:uncomment_begin(google-only)
+#
+# py_proto_library(
+#     name = "test_message_proto_py_pb2",
+#     deps = [
+#         ":test_message_proto",
+#     ],
+# )
+#
+# py_binary(
+#     name = "split_gen",
+#     srcs = ["split_gen.py"],
+#     strict_deps = True,
+#     deps = [
+#         ":test_message_proto_py_pb2",
+#         "//tensorflow/python/lib/io:file_io",
+#         "//tensorflow/tools/proto_splitter:chunk_proto_py_pb2",
+#         "//tensorflow/tools/proto_splitter:split",
+#         "//tensorflow/tools/proto_splitter:util",
+#         "@absl_py//absl:app",
+#         "@absl_py//absl/flags",
+#         "@absl_py//absl/logging",
+#     ],
+# )
+#
+# py_binary(
+#     name = "split_graph_def_gen",
+#     srcs = ["split_graph_def_gen.py"],
+#     strict_deps = True,
+#     deps = [
+#         "//tensorflow/core:protos_all_py",
+#         "//tensorflow/python/lib/io:file_io",
+#         "//tensorflow/tools/proto_splitter:constants",
+#         "//tensorflow/tools/proto_splitter:split_graph_def",
+#         "//tensorflow/tools/proto_splitter/python:test_util",
+#         "@absl_py//absl:app",
+#         "@absl_py//absl/flags",
+#         "@absl_py//absl/logging",
+#     ],
+# )
+#
+# py_binary(
+#     name = "split_saved_model_gen",
+#     srcs = ["split_saved_model_gen.py"],
+#     strict_deps = True,
+#     deps = [
+#         "//tensorflow/core:protos_all_py",
+#         "//tensorflow/python/lib/io:file_io",
+#         "//tensorflow/tools/proto_splitter:constants",
+#         "//tensorflow/tools/proto_splitter/python:saved_model",
+#         "//tensorflow/tools/proto_splitter/python:test_util",
+#         "@absl_py//absl:app",
+#         "@absl_py//absl/flags",
+#         "@absl_py//absl/logging",
+#     ],
+# )
+#
+# py_binary(
+#     name = "many_field_gen",
+#     srcs = ["many_field_gen.py"],
+#     strict_deps = True,
+#     deps = [
+#         ":test_message_proto_py_pb2",
+#         "//tensorflow/python/lib/io:file_io",
+#         "//tensorflow/tools/proto_splitter:split",
+#         "@absl_py//absl:app",
+#         "@absl_py//absl/flags",
+#     ],
+# )
+#
+# copybara:uncomment_end

tensorflow/tools/proto_splitter/testdata/bf-split-tree.pbtxt

@@ -0,0 +1,100 @@
+chunk_index: 0
+chunked_fields {
+  field_tag {
+    field: 2
+  }
+  field_tag {
+    index: 0
+  }
+  message {
+    chunk_index: 1
+    chunked_fields {
+      field_tag {
+        field: 2
+      }
+      field_tag {
+        index: 0
+      }
+      message {
+        chunk_index: 2
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 0
+          }
+          message {
+            chunk_index: 4
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 1
+          }
+          message {
+            chunk_index: 5
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 2
+          }
+          message {
+            chunk_index: 6
+          }
+        }
+      }
+    }
+    chunked_fields {
+      field_tag {
+        field: 2
+      }
+      field_tag {
+        index: 1
+      }
+      message {
+        chunk_index: 3
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 0
+          }
+          message {
+            chunk_index: 7
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 1
+          }
+          message {
+            chunk_index: 8
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 2
+          }
+          message {
+            chunk_index: 9
+          }
+        }
+      }
+    }
+  }
+}

tensorflow/tools/proto_splitter/testdata/df-split-tree.pbtxt

@@ -0,0 +1,100 @@
+chunk_index: 0
+chunked_fields {
+  field_tag {
+    field: 2
+  }
+  field_tag {
+    index: 0
+  }
+  message {
+    chunk_index: 1
+    chunked_fields {
+      field_tag {
+        field: 2
+      }
+      field_tag {
+        index: 0
+      }
+      message {
+        chunk_index: 2
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 0
+          }
+          message {
+            chunk_index: 3
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 1
+          }
+          message {
+            chunk_index: 4
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 2
+          }
+          message {
+            chunk_index: 5
+          }
+        }
+      }
+    }
+    chunked_fields {
+      field_tag {
+        field: 2
+      }
+      field_tag {
+        index: 1
+      }
+      message {
+        chunk_index: 6
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 0
+          }
+          message {
+            chunk_index: 7
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 1
+          }
+          message {
+            chunk_index: 8
+          }
+        }
+        chunked_fields {
+          field_tag {
+            field: 2
+          }
+          field_tag {
+            index: 2
+          }
+          message {
+            chunk_index: 9
+          }
+        }
+      }
+    }
+  }
+}

tensorflow/tools/proto_splitter/testdata/function-large-nodes.pb

@@ -0,0 +1,17 @@
+÷
+ô
+
+fn3
+Const_0Const*!
+valueB
"C3D?7�	>ÿt²>_
+Const_1Const*M
+valueDBB
"8>�h>ñÞ²>wd>|(˜<8Ú>¨éL?}\°=¾Ãu>»œà>Ûm€>D•”>]sE?1,?
+ü>_
+Const_2Const*M
+valueDBB
"8”–×>R\â>7Ö<OwE>cý	?KÐÖ=¥?î�?f\?ˆð8?ˆ&°<õZ²>&èö>×õ¤>_
+Const_3Const*M
+valueDBB
"8=‡?ÙÅÐ=_"?JlG?N4 ?oôJ?EF}?Fã=Ov?N;m>Â\?!(>?O7?Kº–>3
+Const_4Const*!
+valueB
"K÷¹>H?Û+f?_
+Const_5Const*M
+valueDBB
"8=aý>øÈ¡>\G?�b?IY"?ëuG>L±{<9\?V·î= ñ®=v.x?G“>òc?ì…z?

tensorflow/tools/proto_splitter/testdata/many-field.pbtxt

@@ -0,0 +1,31 @@
+field_one {
+  repeated_field {
+  }
+  repeated_field {
+    string_field: "inner_inner_string"
+    map_field_uint32 {
+      key: 324
+      value: "map_value_324"
+    }
+    map_field_uint32 {
+      key: 543
+      value: "map_value_543"
+    }
+  }
+}
+map_field_int64 {
+  key: -1345
+  value: "map_value_-1345"
+}
+nested_map_bool {
+  key: false
+  value {
+    string_field: "string_false"
+  }
+}
+nested_map_bool {
+  key: true
+  value {
+    string_field: "string_true"
+  }
+}

tensorflow/tools/proto_splitter/testdata/many_field_gen.py

@@ -0,0 +1,102 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+r"""Generates ManyField test data for Merger.
+
+Constructs chunked proto test data containing various field types for
+Merger::Read and Merger::Merge.
+
+Usage: bazel run tensorflow/tools/proto_splitter/testdata:many_field_gen -- \
+    --path=/tmp/many_field
+"""
+
+from collections.abc import Sequence
+import os
+
+from absl import app
+from absl import flags
+
+from tensorflow.python.lib.io import file_io
+from tensorflow.tools.proto_splitter import split
+from tensorflow.tools.proto_splitter.testdata import test_message_pb2
+
+# Example path: /tmp/many_field
+SPLITTER_TESTDATA_PATH = flags.DEFINE_string(
+    "path", None, help="Path to testdata directory."
+)
+
+
+class ManyFieldSplitter(split.ComposableSplitter):
+  """Splitter for ManyField proto."""
+
+  def build_chunks(self):
+    self.add_chunk(
+        self._proto.field_one,
+        [
+            test_message_pb2.ManyFields.DESCRIPTOR.fields_by_name[
+                "field_one"
+            ].number
+        ],
+    )
+    self._proto.ClearField("field_one")
+    for map_key, map_value in self._proto.nested_map_bool.items():
+      self.add_chunk(
+          map_value,
+          [
+              test_message_pb2.ManyFields.DESCRIPTOR.fields_by_name[
+                  "nested_map_bool"
+              ].number,
+              map_key,
+          ],
+      )
+    self._proto.ClearField("nested_map_bool")
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+
+  proto = test_message_pb2.ManyFields(
+      field_one=test_message_pb2.ManyFields(
+          repeated_field=[
+              test_message_pb2.ManyFields(),
+              test_message_pb2.ManyFields(
+                  string_field="inner_inner_string",
+                  map_field_uint32={
+                      324: "map_value_324",
+                      543: "map_value_543",
+                  },
+              ),
+          ]
+      ),
+      map_field_int64={
+          -1345: "map_value_-1345",
+      },
+      nested_map_bool={
+          True: test_message_pb2.ManyFields(string_field="string_true"),
+          False: test_message_pb2.ManyFields(string_field="string_false"),
+      },
+  )
+  file_io.write_string_to_file(
+      os.path.join(SPLITTER_TESTDATA_PATH.value, "many-field.pbtxt"), str(proto)
+  )
+
+  ManyFieldSplitter(proto).write(
+      os.path.join(SPLITTER_TESTDATA_PATH.value, "many-field")
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

tensorflow/tools/proto_splitter/testdata/split-large-constant.pbtxt

@@ -0,0 +1,90 @@
+node {
+  name: "Const_0"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 3
+          }
+        }
+        tensor_content: "/O5?\342\020\271>\215\326{?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_1"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 3
+          }
+        }
+        tensor_content: "\327\2177>\024)G?\007\027\224>"
+      }
+    }
+  }
+}
+node {
+  name: "Const_2"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 241
+          }
+        }
+        tensor_content: ">!\343>R\362[?\025a\003>u\224\007>9\237\n>R\264b?\330\310k?c\350\315=^\335Q?\332\254s?\"$a>\366\'M?4Yz>\363\036\017>\270\363U?\260\300\351>\350z+?\270\3270?\213\3651?>\024O?\247:K?W\"\211>\021\001P?v\346\262>=\244\307>*\t^?\336=O=\2265\206>\241\322\023?\'\3658?\256p\014?\344\253V?N\247o?j\013\311>\r\267??t\365*>\245\030\007=\334\300]?\005\221\'?\032W\266>Q\033\005?\375p\341>\324B=?\360 C?\375\317T?\016\001\265>\007\347\357>\325k\235=\004]}?\266\212\253>f\266C?\326\032\"?|\032\371>\311\016\252=\240\355d?\340D~?\310\343\'>\237\305\023>P\035;?\327\003\223>\246\202w?\320\026\'?y\351)>\335J\300>\221\362l?\r>\032?\3556\234<\231\017\253>\274\326~?\316{Y>N/\304>2\274l?\272\332\255>\254\373\032?)\334!?#jX?\360\336&?N\263\010?T\251\037?\240\014\344=\371\\\371>\261X\270>M\257\026?\0278\365>]\200\220>\377\200\200>[\267t?\220\353\013>8>~?\276\035\020?\033\217E?\004\200}?\223\373>?n\320\244>\274T\010?\347\363\260>\000|\010=u[\264>\355AS?\216gV?\231\341r?\002\360\033?g\204\027?E\354??\025\246t>`\025\000?\312\035\206>\335}x?\305\206\253>\275\356f?\306\251S?i&\206>:U\313>\313\354^>\231\353\t>\035S4?m\376\323>\277\337\372>EXx?1\216\265>\310\220\003?\272\021\037?\010g\305>\304\270U>\270._>F\333D?\343\357\005>\250\030\310>\2521\"?#6~?\305\002\210>)\3320:f\025\t?V\022v>\321~\214>.v\017?\204p1?\363\240\216>\357\225\221>\n\332\234=\025\0033:\272F\264>!Mn>\335f\350>\243\334/?\023\236y>Q\220\"?\316st?I\023\355>\0002\331=\377H]?\364\010d?}\326(?\301\251)?|\354\016?K\212B?\002\237t?\200\"M=\250\016&>|\243\360>\372p\177?\031xC>\323\177\310>\323\215\226>\230k|?~\227\n;\344E\004?\252\357)>\257\353*?\\\325\'>\255\365\214=\022Q(>G}[?f@\177?y\310\032?\010o\267=%\362\326>4\361\341>\026n\020>\375RV?^\203\023?\222\336x>I\276P?\020\212\304>\342;\350>b\037\330>s\236\027?\306\346\336>\322z\030?m\247d?\320zB>c\351N?\001\365\r?\336g\027?\253\031\037?\333\235\\?\000\314E?Q\302y?\206K\243=r^\'?\201\236\243>@Kj?:\001\326>\330L\266>\322Yj?9D\261=\357\216w?c\276C?\201\204\326=\235-E>\325Z\254>76_>S\304m?\346*\250>\241\017D?a\301\014>\332\316\372>\230\007?=\241\'b?.\357s?N\350\022<\032p+?\345\256d?\023\231\023>I\257\017?\\~V>\'\327]>\377A\'?GGL?\371\372f?\323\341H?pB\207>\344\322\214>r\271(?\363R\265>%Rb?G\1770?X\351C?\313\223\261>i\255\204>\255\n\013?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_3"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 3
+          }
+        }
+        tensor_content: ":\267\364>0\310\222>\260\224`?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_4"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 241
+          }
+        }
+        tensor_content: "-\020\353>\212!k?.\264\007=r\342\323>=\023P?#\356\025>o\310q?\352\237x>\003\264p>\201\212\004?\262#J>\222\005\"=\304\272%?\021\\\362>\205\256z?d\3042?\"~\263>\300\303\342=\332\003*?\356\301O?\303\316l?\0073\336>\334\033\213>\013\371\033?\275\033\311>+\204P?\221\000\022?5.t:id\005?B\265\245>k\'\202>\203h\304>9\236\234>.\202@?-\303U?\215\0326?\270\346\005=ro\205>\242zu?\253m\342>_\234_=A\237\303<\3610H?\211\272\010?\235\307\216>\"?g?*\005\377=F\231\330>\2325\001?\221\006_=\025\2762?\206wZ?\211\265\342>\242C\007?\237T^>\306\211.>\2010\265=\370yf?hw\347<\357-)?8+b?\253\246\325>2d\222>\242\246\004?\031\022\260:z\030\365=K.\266>\302\224G?\311\330*?\325\301\307>\257\201\271>\242\216\355>D7\024?L\275S?\254\224\250=\014i\250>p)=?\264i\346>dGS?@\\v?\214J\310>\256\2523?\226\245\005?\357)R?(:d?T\265\232>/ZV?\214a4?\327\242V= \376K?bp\376>v\236%?\242\206??\"\276\257>\271\014+?\0224\002>\365\010\313>\214k\000?\236\265\023?\037UG?\237\274B?B`\372>\371u\033?\247x\324>\266\346\321=N\024f?\254\205\310>\203n\030>:\312h?\315\355\212>\266\244\026?\253r\375>[.\242=B\225\277<\314\202b>wN\374<\032\020\014?\300i\351>\245\372\275>O\242O=5fE>\t\333\316>\021^\020?e\270\305<\255\277k?\204Z\227;$2\337>@V\004?\371\375\\?\007\323\004>@\206\177?\205\203\346>l\036\025?\332\264\207>4#\307>[\3119?\344\342\024?C\177(?\305Sk?\377wl?\347\351P?\004\t\354>I\214\264=`\227\366=H=\031?c\251P?Q\305\030<\336\036\006>\326\214\311>{k\024>\032\333\263>\023\305\350=~$\314>\237\225\005?+vi?\027\267\314>\014\223P?l\315\375>\265OZ?\257\036\310>\014P1?\250\001\037>;\366\030?\005~M?mH<?\325\257\327>\352\244k>\311\217(?\243\305\367>$(>?`\234\235>\322[\n?\227\237h?\350Vd?E?\222>\244\377\226=\362\205\235=-\342W?\211u\001?\276V1?_\370\365>\361M\250>\004\335M?2Y2?\202\273.?#X\256>\245\344!?H\2061?\027\340\372>\325\311|?x\n\353>\217\307\310=\242\325E?B(\216=x\223;?D/\241>b\2407?i\024y>\311{`=\2037\354=k\342\220>\214\340H?\001\233\263> \021-?\345x$?m\242\211>\345i+?\305\276)>\224\266`=\377.3>\213\251m?\330\031\337>JTv?\343\264\361>\340\332J?{\265\245>\357v\334>-\327u?\225\013h?\316\224Y=\274`l?\244>\034?\264A\022?\263\203\257>\016hV=\020]\022?\302P\363>\"*\324>M\343m?^\016x?\204\362\331>\273\354:?\271\320.?\231\2100>*\272\003>\024i[?\262\347\364>Hq2?\024\247\306=\234\027\372>\321\3358?"
+      }
+    }
+  }
+}

tensorflow/tools/proto_splitter/testdata/split-large-nodes.pbtxt

@@ -0,0 +1,108 @@
+node {
+  name: "Const_0"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 3
+          }
+        }
+        tensor_content: "\355,W?y\232{?\234G\234>"
+      }
+    }
+  }
+}
+node {
+  name: "Const_1"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\222\241\363>c5r>>\244P?\345&\324>^i ?\205\'\242>2*\221>\014\305z>a\243@?\354Pj?ZH\210>\215\225\001>d\203/?\255\372k?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_2"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "4\322\026?\2046K?\265\367]>\3672\305>\272\345w?\003\323\370=x\000\246>\301\255\027?\215)6?;\302W?\212#o?\0376\004?\335\356\016>\260Va?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_3"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "v\334H?b\366M?\345\200\021?H\364\247>v\0231?\311\203\036?H\2777?K\227R?\351VI?\343\240\256>\341\333@?\364\251\244>lnY?\201\336d?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_4"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 3
+          }
+        }
+        tensor_content: "n\344T?2\220m?v\001\252="
+      }
+    }
+  }
+}
+node {
+  name: "Const_5"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\200\216\007?\377\350o?\337\363\316>\275e\337=\313\263\t=\326Z\001?\263lI>\322\021_?\377\340\007?\354\321p>}\264\222>\036\177\037?\001\234#>\021\243<?"
+      }
+    }
+  }
+}

tensorflow/tools/proto_splitter/testdata/split-lots-nodes.pbtxt

@@ -0,0 +1,270 @@
+node {
+  name: "Const_0"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\016\300\252>\\\025$?\205M\365>~>\036?X\237\016=\232=]?Y\002*? \263\216>V\002x=\372C\355=\243\356\376>\213M\023;_fv?\'I~>"
+      }
+    }
+  }
+}
+node {
+  name: "Const_1"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\235\325\005?OO\374>m\206\004?0\244\320<\201\036\317>\304\264X?\320Yb?\350\313\013?\202\241\352=d\215*?\306\336^?\3438y?W\037~?K\214h?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_2"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\342&\035?\246\343\250>\324k\236>E\243Q?\253/\026?+!e?Q\353\302>\203mc?^7\311> \213[?\222\004??H4\277=\3065t?\374\346N>"
+      }
+    }
+  }
+}
+node {
+  name: "Const_3"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\270\255M?\033\203\312>\225\263`?%\'\350>\177P+?\2633\030? \2630?\262\243(>>\221\373>\203\203|>\314\302\301=\026!\227=\2329 ?\t\255\326="
+      }
+    }
+  }
+}
+node {
+  name: "Const_4"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\320\333\265;\276r.?P\271\340=\302?\035?\314\2163?j\n{?#\216\255>\256\r\374>7\266\337>\024+4>yH\315=\244vW?\354\007\262<p%=?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_5"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\375\233\202=\2257\254>8\337\004=\234\227h>I\327j?\177S$?L\375\005?G\306s?z\2479>\257\330,>~,\037?m\313R>}\210\r>\330\224w?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_6"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\360\020\353>\351\206(?\360f\244>\033\257\034?\377\321>?J[\335>\363\372\036?1\336e>\263\207\232> [\010?\312d\274>\351\317v>\031\320\320>\377}\017?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_7"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: ">\242\232=4/\036?\250\203\344>\010-0?\346\022\212>\211\022(?w{\025>:\025c?\035\374\246>\374\255[?\035_\246>fj^>\346\010\225>\230\3539="
+      }
+    }
+  }
+}
+node {
+  name: "Const_8"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: ">\211n?pZ\001?*\3054>\276\347W?\346\374\251=\251\257\350=\003\245\362>\nvs>\360\251g?\2039(?\235\202\370>L\301>?\373J\034?\026\225\225>"
+      }
+    }
+  }
+}
+node {
+  name: "Const_9"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\222\342\226>pN\212>\256c7?NS\263>\001\312\225>y0`?w/\301>3\3369>\312,\243=\353\036\030?\267\022\351>\314\246l?\004O\245>\303\035\322>"
+      }
+    }
+  }
+}
+node {
+  name: "Const_10"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\363\030R?\215\367j>I\364\301>8\331\022?0e\321>\007\2240?z\275|?\327u\250>z\225U?t\301{?\223$\010>0\006\025?\267\033\n>\336e\351="
+      }
+    }
+  }
+}
+node {
+  name: "Const_11"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\213\336\315>d\t5?\364\331\010?\267G\205<\277\002\204=\036\343j>\017\267\274>I\275F>\005\027\222>\\\226 =\254\267m>\245\373!>=\\\014?\307\314\240<"
+      }
+    }
+  }
+}
+node {
+  name: "Const_12"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "\226rG?\232\247Y>A\027n?.\321H?B\3638?B\244*?\340\\\030?\227[7?\036\375\032>\nv$?)\313\227>\\\255l?\033G:?,\r<>"
+      }
+    }
+  }
+}
+node {
+  name: "Const_13"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "a\332\223=\006z\'?\n\225\021=\263&\013?\355\222\267>\0307\026?\356\200\270=\227\361e?1(g?\237\032\272;\2447\252>\213a8?)\276\"?]!\021?"
+      }
+    }
+  }
+}
+node {
+  name: "Const_14"
+  op: "Const"
+  attr {
+    key: "value"
+    value {
+      tensor {
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim {
+            size: 14
+          }
+        }
+        tensor_content: "v\261k=\213S\216>\346\257f?!\247r?X\330\314>\356\001~>\005\177`<\316C\">\357\250\032?\020h\027=\260\002\007?\261\2609?Q+\216<\323+\t>"
+      }
+    }
+  }
+}

tensorflow/tools/proto_splitter/testdata/split-standard.pbtxt

@@ -0,0 +1,332 @@
+meta_graphs {
+  graph_def {
+    node {
+      name: "Const_0"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 16
+              }
+            }
+            tensor_content: "\235d\331>N\217`?M!\241>\023\'|?P\006$?\017\225*?\332\016\265>\311\221\233=2\205\370>\303\216i?Y\321\271>\221\032\233>\343\335&?;\031\305>&=\317>\360\220T?"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_1"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 16
+              }
+            }
+            tensor_content: "{\260\007>%\n1?\200\205\273> \306N?\337\201\363>B\2348?C\366\306>A{e?g2}?\225_I?\356@!?^G\265<&\013\260>Ca\021?\300\355k>\363y\361>"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_2"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 241
+              }
+            }
+            tensor_content: "*\272+>\203#i>6r[?\254\206/>l\002S?\242\213K?K\214\205>\023\275\312>\256EM=y\321\n?\014\274^>\370\037{?\367\016t?\274\327\241>[\376\277>\225\036}?x:]?\272tO?\272\364\032?9,;?`I(?\377t\222=~;\355>o:\356>\225<\022?t\350\266>\000:O?&>\201>n\230\273>n\266\307>\370\034\260>\301\375n?\230\357\013?\346\312\021?\274D\264>\357\016\237>m\337\013>\373\3025=\246\262\317>\315\351z?\001\376#?\001\370\374>\305\277^?\031\231\342>\307\266\r?\364\236\032?\000\r\211>\223\210v?bq\307>\334\341\264=\006\274\265>\310\232\017?\013D\372>[\004\014>\032!^?R\212\301>\227\"\324>\024\311.?(\371\375>\207\007#?\031#k=BR=?\030Bj?\001\215)=e\351K?s\376\350>\330)\247;\337\030]?|d\275>\276\307\273=?\372r?u\262\335>\353\205z?\322\362\305>\313y*?r\316\270>Zf\325>\377\341\023?}+\234>\023\265\000?8>A=-\302\n?\323s\352=\331\371\215>G\000\247>`\001L>_\004F?\322\317-?\003dQ?h\007\005?o@S=\025\221c?U\203\245;\270\223d?\331\364\001?\3248V?\207\334H>\r\245\342>\346\331\311>\244\024\274=\300\023m>D++?\346\237M?w\020\304>\355$}?<A\032>\371h\333>\344\274|?\233\247M>\277\377\217=\221\031\275>V\214\234>\222\347:?\203\001\007?\312\241\025>\274\273(?\342\344\035?\311.\177?\346\246\312>\267\205\304>\370/H>k\213`?\0041\313>\234\333\264>^\257\032>\023<\022>[\323\026?\230\254S>\233\360B<p\321\002?b\346\341>9\244P>\313\351\226>\217\365\020?E(`>\254%A?\320q@?M_<?\363\037\010<\200\203d?\036h\212>\306\264\036?g\206\241=\327\341\036=P\033\366>\266\032\266>\254\224X?\361:a?#\017T?\017\220\016?\r\252\261=.\366\245>*r\303>\031\370\222>@T<?\027\235\021?\0214T?8-.?`~4?\375;j?u?\266>\267\215\207=\031\323{?\377y\'>\001\240f?Iv\276>s\214u>\356\354K>\315\256z?\210+\361=\250\372\254>\346\246\253>\223R\016?\201bR?t\227\215=\335j}?V\365\211>I\"]?~\325\354>\213\275*?:\3371?\263\267:?8\033\321>\352C\013> 7\023?\255\214\336>\362\244\243=\360.\355=1Z6?\201\177L?\332\277<>\273\264\010=\036A\226>\277O\254>\307\024G?\347\030,?\355b\327=\257SV?\037\336\017?\262\232\001?\034<k?\'\266\006?\252Wv?\013`\027?B\n/>\210\237\";|(\021>\022\023\336>\362\205?>M\220m>m\304\364=z\030\220>|\375g?\335\013h?\"\321\215>s-x?\342\334\036?\244\037\367>\334\245L?,\317d?\301\317r?\335\371\265=v\331\025?\035qT?a\253\004>6&w>\3238q?J\252\321>h2O?\316ra?\"G\206>\013\275\003?u\364\n?9:J?\t\305K?3\023\365>\241-\037?\225\354~>>\264\020?C\2066?\333\3446?"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_3"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 16
+              }
+            }
+            tensor_content: "#\375.>\021\223\021?\313\367K?\327.\033?w\206m?\273q\027?\250\213I?\201mV>\266\377b>\246\212\003?h\004]?\214H\010>\020\274\001>\033\035\010>\234\373x>r\312i?"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_4"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 241
+              }
+            }
+            tensor_content: "c\203;?\335\037C>\007\316E>\245\010\332=k\344\306>\201r\025>\330\274\020?h\272t?n\001\\?\264\226\314>\014z\014?\233\014J?KAU?\327\225V?[\007\240>\005\277\253>\341\250 ?5\362^?\355\366\246>\324&h?)\211\005?\343\371\347>\347sb?\234\236\325>E\004\037?\rkz=S<_>+\350\277=$\212M?\211S{?\222q1?q\352\233=By\245>\"r\354=\214\345\236>\177}\037?y\024\306>Ki\037>\021\216\007>\227\204\026?\326\371\022?\352\364~>\221\215\272=\026\372x?\031%c?/\336D>\277\363*>mH\r=G\222\014?4P&?\250AR<\"\307g?^\303{?\312\350\340>\270\234\264=\323\000\354>\0050+>\244\206(?\341\000^=\270\320\226>z\365\221>\014\007\347>\277^\030>w\035\340>G=*?\310c\236>2\037%=\361Qq?\215\261j?\033\204\345={T\217>\253N\333=\211]\004:\t\301\221>\357\031\016>eG\272>\"\370L?&\004j?\230\275\275<\300=5?\366\036a>\'\254X>\272zV?fY\220=\231\025w?\322\037&>m\220\335>\177\300\372=\210B\010>\326\313p?\305\340\325=\362\343k?\227\252`?\372Z\326>P9Q?\004\317X?\363l\013=\233p\177?\342T\021?f\224\365>\370\270A?\366\327\220=\343[\016?!\2033?\205\344\320=\336\0007?\360\n\t?\2012\000?\327(B>\007\252Z?\253\202\367>\225\356`>Z\2026?|J/?\311>:>r\023\006>\350qf?\265\317@?\243\017p?\016)T=\324\374\010=g\341J>)\243y>\257\004\252=\223\277|?1\237u? <%?\313|\\>\2164\216=\206\0222?XY,?\336b$?@\033G?\313\206\255>\315\372(>B\345e?\227\3123?\310\231\251=\1772\036?l%|?\312\032\334<\253Z\221>\357\020<?\256\201-?\220\342(?\035\r\324=\025a\326>\270\2779?C1S?\\\007\003>\374E\363>\222\223<?P\334\"?\204\177%?\350\207\371>)\356E=\37321>\240\220}>\356\000A?}\366*>\264Of?\030lm?\372\224\r?\351\211\252>L\031|?\215\237d?\337yb>\25277?mn??S\267\231>:\317]?\364\350\334>\324\247\007?\232] ?\n\026\311>\323\024E?r\372%?\017\362E>\250\214\317>\353l->\337IT?_\210o?o\037\361>$ \032?zxV>\212\327\037?W\317\177?o\031\006?\331\261\036?V\254y>\333\300P?\315\302\000?6\312\235>\034\373\304>\2003 =\344Oo?\210\030 >\265\\\353<P\350r?g\336\336<\342f\325>\253!\207>py\023?u\014\036?o\017<?\271\267T?\217\234\314>\243@\323>\241=\373>\212)\225=\255\234,?o\354\213=\204\235S?Y5c?>\310\005?\014*\024?\304\331;?\002|\226>\231\214\376>\341S\214=g4N?)\272\303>\225\036P>\213\324\\?\353i\341>\241\230j?\302\367m=\256\215R?\360\231\226>\2042\321>\025{\311>\332@\265>$E\017?\216\323j?\260\365\263=\276\353\276>\207\272r?g\315\212=\303\302\335>\347\026\376=\254@;?"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_5"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 116
+              }
+            }
+            tensor_content: "\036\014n?H/\027>\375be?\247\2065?v\303\366>\266\357\025?\266G#?\364c\304>\245j\035>{\273 ?\312\030~?\017o\031?/\345L>\272\254^>QQ\005??\366\201>\217N.?\274\013j?\0300{=\243\'0?\343=\325>\237v\010>;UQ?\322E\002?G^<?\tg4=\271\345^?y,\363>\204\206\251>\036+!=\275/S?\005\002\260=\203\3374>\\\300G?y\204\177?\257\254\007?\235\230+>\264\365\003?\014\261\261>\025\322\026?\2375\225>\031\263T?\313N\033>\237\325\016?ZtH=\204\331>?\215\343\010?\253\354L>\360\325>?5\224;?\232D\256;?\031B?\337kR?\3677s?`0\367>\332@\n?\024\245a?\352g\373=\303l7?\224\001\313>n\364\025?\006c^?#N\022<zs7?fD\007?$\021W?\"\027!>\233]\207>\315\356v?)\2075>\217}L?\345H\364>\333\033\374;p\271\034?\027Ke>@l\n?=\235\225>\t\264L?\256\343b>u\203\325>\234}\273=\340\\\333>\370k\245>\0021\n?\332\027\330>X\336S=7N\"<\227\335O>\375\025\037?\213Ee?\244\262\036?\210r*?\343zU?\022\347\025?\021rf?f\037a?\354\240\353=\216z#?Vr\256>\016\033&>\301\252\307>C\271\311>`\320\253>\307p\235>K\302f>\362s\272>\376~\377>\204Ct>\203\t\036?VjO>\023{&?Gf\036?j\017*?6pw>\244\272r?\032\"\021?"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_6"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 16
+              }
+            }
+            tensor_content: "\212\220q>\213\277y>4?\006<1\226\177>o}\271>b\374\016?\237\366,?\357h->P\312\230>\010\033_?\232\240\013?\004j\206>\023\021\007?%\024Q?\277\3634?\004\317*?"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_7"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 16
+              }
+            }
+            tensor_content: "\206\226v?4\211F?c\250\353>\'G\t?:\241%?\272y|=\020o\n>NN\265>9\3301>\313\377\325>&\302\013>3\035%?x\217\017?\037\255\345>P(\334>g\263<?"
+          }
+        }
+      }
+    }
+    node {
+      name: "Const_8"
+      op: "Const"
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_FLOAT
+            tensor_shape {
+              dim {
+                size: 16
+              }
+            }
+            tensor_content: ".\356\201>\t\246\371>\234\014p?D\3626?j;\013=V\232h>\032\327w>I\000B?\320\336Z?(\364\234>\020$\223=~q\211>v\345&?a\257\000?\030\221\360>\350pX?"
+          }
+        }
+      }
+    }
+    library {
+      function {
+        signature {
+          name: "fn1"
+        }
+        node_def {
+          name: "Const_9"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 16
+                  }
+                }
+                tensor_content: "\021\"/?\032\335\354;\210\352Q>!\344\341>\243[\204>\245\227h>\020\177\006?^\357\203;\273\305T?\257\213\370=p\352\230>\261\n\220>\256\375m=KO\267>\\|w>\241\324q?"
+              }
+            }
+          }
+        }
+        node_def {
+          name: "Const_10"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 15
+                  }
+                }
+                tensor_content: "CT8?\254\211Q?HHd?\240\210\351>\352BJ?\013\253\265=\0305\357=&\225\016>\362\267\211>?G\332>*\3031?\356\315\255;Z\344$?\023$X=\344\272\271>"
+              }
+            }
+          }
+        }
+        node_def {
+          name: "Const_11"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 15
+                  }
+                }
+                tensor_content: "\317\350l?A\030\357=b\313\246>\304@*>]~\205=t6\010>\251]\000>\205ms=\236\252\006>/\343\236;T\226\364=/b\273=v|E?\374\375\014?\303\317\027?"
+              }
+            }
+          }
+        }
+      }
+      function {
+        signature {
+          name: "fn2"
+        }
+        node_def {
+          name: "Const_12"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 15
+                  }
+                }
+                tensor_content: "\205\324l?\274|H?)\235!?\242[??\3551\307>\355\2040?%\310\342>\034\025E?\302j\'?\237Q\350>*\r^=a\317;?p\271=?\272MA?\267&\346>"
+              }
+            }
+          }
+        }
+        node_def {
+          name: "Const_13"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 115
+                  }
+                }
+                tensor_content: "N\265\276>.X\304>\352%k?\225!\245>\305\304\235>\346\211{?\3272*?R\321f?\320\300i?Iy@>\271x\203>\002\322\274=ioI?\t\225\372>\362Bp?\2031W<\022\337\002>E3\361>\366\214\274<\2274\036?\337\220\371>,B\336>*\361C?\250[\263>\236\t\020?\346Y\261>\336\232U>v=\325>f\277f<\207C2?h#=?_\215h?\314\211\215>\032rn?\214CH>\'\232\342=\204\376r=\374\237\'?x\310\265>\005\223\315>\260\270\016?\2179\255>\370\t\025>\224RE?ay9?\004\372=?\237\024G?!\034\351=\0231e?M\327\234>`\357f?\375\355t?\365\274\177?n\220E?sN\267>\253tQ?\303\277K?h\351\272>\031\277w?\007\341.>\240\024S=\217\177y?\271z\010?\323\231\356>\306k\034?\031@\264>\024\371\031>\036\211\035>W\276K?\001\272K>\223\'v?\213\212\376>\372F$?6\276R?\257vb?V\341\202>\004\023\032?=P\371>W\010F?\007\222\376>]\232\321>\231\252\370>\'7.>,\233k?`_\376>I\337d?\363\220b?hS]?\203\377u?\177Y\007?\023\177\306=\234\2534?\322\257\262>\225R\010<\371 m>4\334P?\337\241\345>E\001\271>\352\226O?*|@?\203\352a??#\010?\277\243\005?]\324\361>3\007;?\"\3753>\3427\000?ZW\022>\0037G?\350\220\273>XZ\277<q\252m?\342\315\311>xJd?\n\343\035?"
+              }
+            }
+          }
+        }
+      }
+      function {
+        signature {
+          name: "fn3"
+        }
+        node_def {
+          name: "Const_14"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 15
+                  }
+                }
+                tensor_content: "\367\356)?\020y\310=o\2048?\243 >?\312\203T?a\273W>n\212m?\232C6?\370\331/?/\374\253>\372R ?\257|\017?\230\256X=\204\330\"?\0007\366>"
+              }
+            }
+          }
+        }
+      }
+      function {
+        signature {
+          name: "fn4"
+        }
+        node_def {
+          name: "Const_15"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 15
+                  }
+                }
+                tensor_content: "&\036\000?\271\260b?\375\307h?\277u\371=Vl[?rP\321>:\306\010?\274\001\002?\267\\k?\374e\006?\347\355;>RlC?\017\370\330>\362\002\276=\336p\357>"
+              }
+            }
+          }
+        }
+        node_def {
+          name: "Const_16"
+          op: "Const"
+          attr {
+            key: "value"
+            value {
+              tensor {
+                dtype: DT_FLOAT
+                tensor_shape {
+                  dim {
+                    size: 15
+                  }
+                }
+                tensor_content: "J\241F>\232\"i>H(z?b\227\274>t\331\233>\224\361O>&\210\210>au\316>\231\210V?\224\241j>\377P&>T>\303=C\371 ?O\232;>\216\016|>"
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}

tensorflow/tools/proto_splitter/testdata/split-tree.pbtxt

@@ -0,0 +1,28 @@
+val: "0"
+child_nodes {
+  val: "010"
+  child_nodes {
+    val: "01020"
+    child_nodes {
+      val: "0102030"
+    }
+    child_nodes {
+      val: "0102031"
+    }
+    child_nodes {
+      val: "0102032"
+    }
+  }
+  child_nodes {
+    val: "01021"
+    child_nodes {
+      val: "0102130"
+    }
+    child_nodes {
+      val: "0102131"
+    }
+    child_nodes {
+      val: "0102132"
+    }
+  }
+}

tensorflow/tools/proto_splitter/testdata/split_gen.py

@@ -0,0 +1,179 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+"""Generates test data for Merger.
+
+Constructs depth- and breadth-first tree-like chunked protos test data for
+Merger::Read and Merger::Merge.
+"""
+
+from collections.abc import Sequence
+import os
+from typing import Optional, Union
+
+from absl import app
+from absl import flags
+from absl import logging
+
+from google.protobuf import message
+from tensorflow.python.lib.io import file_io
+from tensorflow.tools.proto_splitter import chunk_pb2
+from tensorflow.tools.proto_splitter import split
+from tensorflow.tools.proto_splitter import util
+from tensorflow.tools.proto_splitter.testdata import test_message_pb2
+
+SPLITTER_TESTDATA_PATH = flags.DEFINE_string(
+    "path", None, help="Path to testdata directory.")
+
+_CHILD_NODES_FIELD_TAG = (
+    test_message_pb2.StringNode.DESCRIPTOR.fields_by_name[
+        "child_nodes"
+    ].number
+)
+
+
+class StringNodeSplitter(split.ComposableSplitter):
+  """Splits a StringNode proto with N strings into a tree with depth N."""
+
+  def __init__(self, proto: test_message_pb2.StringNode,
+               chunked_message: Optional[chunk_pb2.ChunkedMessage] = None,
+               **kwargs):
+    super().__init__(proto, **kwargs)
+    self._chunked_message = self._chunked_message or chunked_message
+
+  def add_chunk(
+      self, chunk: Union[message.Message, bytes], field_tags: util.FieldTypes
+  ) -> None:
+    """Adds a new chunk and updates the ChunkedMessage proto."""
+    assert self._chunked_message is not None
+    field = self._chunked_message.chunked_fields.add(
+        field_tag=util.get_field_tag(self._proto, field_tags)
+    )
+    field.message.chunk_index = self.total_chunks_len()
+    self.add_root_chunk(chunk)
+
+  def total_chunks_len(self) -> int:
+    """Returns length of chunks stored in root splitter."""
+    if self._parent_splitter is not None:
+      return self._parent_splitter.total_chunks_len()
+    return len(self._chunks)
+
+  def add_root_chunk(self, chunk: Union[message.Message, bytes]) -> None:
+    """Adds chunk to root splitter chunks."""
+    if self._parent_splitter is None:
+      assert self._chunks is not None
+      self._chunks.append(chunk)
+    else:
+      self._parent_splitter.add_root_chunk(chunk)
+
+
+class DFStringNodeSplitter(StringNodeSplitter):
+  """Depth-first string node splitter."""
+
+  def build_chunks(self) -> Sequence[Union[message.Message, bytes]]:
+    if not isinstance(self._proto, test_message_pb2.StringNode):
+      raise TypeError("Can only split TreeString type protos")
+
+    if not self._proto.child_nodes:
+      return
+    for i, node in enumerate(self._proto.child_nodes):
+      self.add_chunk(node, [_CHILD_NODES_FIELD_TAG, i])
+      DFStringNodeSplitter(
+          proto=node,
+          parent_splitter=self,
+          fields_in_parent=[_CHILD_NODES_FIELD_TAG],
+          chunked_message=self._chunked_message.chunked_fields[i].message
+      ).build_chunks()
+
+    self._proto.ClearField("child_nodes")
+    if self._parent_splitter is None:
+      self._chunks.append(self._chunked_message)
+      file_io.write_string_to_file(
+          os.path.join(SPLITTER_TESTDATA_PATH.value, "df-split-tree.pbtxt"),
+          str(self._chunked_message))
+
+    return self._chunks
+
+
+class BFStringNodeSplitter(StringNodeSplitter):
+  """Breadth-first string node splitter."""
+
+  def build_chunks(self) -> Sequence[Union[message.Message, bytes]]:
+    if not isinstance(self._proto, test_message_pb2.StringNode):
+      raise TypeError("Can only split TreeString type protos")
+
+    if not self._proto.child_nodes:
+      return
+    for i, node in enumerate(self._proto.child_nodes):
+      self.add_chunk(node, [_CHILD_NODES_FIELD_TAG, i])
+    for i, node in enumerate(self._proto.child_nodes):
+      BFStringNodeSplitter(
+          proto=node,
+          parent_splitter=self,
+          fields_in_parent=[_CHILD_NODES_FIELD_TAG],
+          chunked_message=self._chunked_message.chunked_fields[i].message
+      ).build_chunks()
+
+    self._proto.ClearField("child_nodes")
+    if self._parent_splitter is None:
+      self._chunks.append(self._chunked_message)
+      file_io.write_string_to_file(
+          os.path.join(SPLITTER_TESTDATA_PATH.value, "bf-split-tree.pbtxt"),
+          str(self._chunked_message))
+
+    return self._chunks
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+
+  if SPLITTER_TESTDATA_PATH.value is None:
+    raise app.UsageError("'path' flag not specified.")
+
+  levels = 4
+  def make_string_tree(
+      string_tree: test_message_pb2.StringNode, level: int = 0, label: str = "0"
+  ) -> test_message_pb2.StringNode:
+    string_tree.val = label
+    if level >= levels-1:
+      return string_tree
+    for i in range(level+1):
+      make_string_tree(string_tree.child_nodes.add(),
+                       level+1, label+str(level+1)+str(i))
+    return string_tree
+
+  def copy_string_tree(string_tree: test_message_pb2.StringNode):
+    new_tree = test_message_pb2.StringNode()
+    new_tree.CopyFrom(string_tree)
+    return new_tree
+
+  string_tree = make_string_tree(test_message_pb2.StringNode())
+  logging.info("StringNode tree generated:\n%s", string_tree)
+  file_io.write_string_to_file(
+      os.path.join(SPLITTER_TESTDATA_PATH.value, "split-tree.pbtxt"),
+      str(string_tree))
+
+  # depth-first chunk ordering
+  DFStringNodeSplitter(copy_string_tree(string_tree)).write(
+      os.path.join(SPLITTER_TESTDATA_PATH.value, "df-split-tree"))
+
+  # breadth-first
+  BFStringNodeSplitter(copy_string_tree(string_tree)).write(
+      os.path.join(SPLITTER_TESTDATA_PATH.value, "bf-split-tree"))
+
+
+if __name__ == "__main__":
+  app.run(main)

tensorflow/tools/proto_splitter/testdata/split_graph_def_gen.py

@@ -0,0 +1,198 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+r"""Generates GraphDef test data for Merger.
+
+Constructs chunked protos test data containing GraphDefs with lots of nodes and
+large nodes for Merger::Read and Merger::Merge.
+
+Example command:
+
+bazel run tensorflow/tools/proto_splitter/testdata:split_graph_def_gen -- \
+  --path /tmp \
+  --graph_type=split-lots-nodes,split-large-nodes,split-large-constant \
+  --export=pb,cpb
+"""
+
+from collections.abc import Sequence
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+
+from tensorflow.core.framework import graph_pb2
+from tensorflow.python.lib.io import file_io
+from tensorflow.tools.proto_splitter import constants
+from tensorflow.tools.proto_splitter import split_graph_def
+from tensorflow.tools.proto_splitter.python import test_util
+
+
+LOTS_NODES_SIZES = [95] * 15
+LARGE_NODES_SIZES = [50, 95, 95, 95, 50, 95]
+LARGE_CONSTANT_SIZES = [50, 50, 1000, 50, 1000]
+
+
+def _split_and_write(
+    path: str,
+    graph_def: graph_pb2.GraphDef,
+    max_size: int,
+    export_files: Sequence[str],
+):
+  """Writes the .pb, .pbtxt and .cpb files for a GraphDef."""
+  constants.debug_set_max_size(max_size)
+
+  if "pbtxt" in export_files:
+    output_path = f"{path}.pbtxt"
+    file_io.write_string_to_file(output_path, str(graph_def))
+    logging.info("  %s written", output_path)
+  if "pb" in export_files:
+    output_path = f"{path}.pb"
+    file_io.write_string_to_file(output_path, graph_def.SerializeToString())
+    logging.info("  %s written", output_path)
+  if "cpb" in export_files:
+    splitter = split_graph_def.GraphDefSplitter(graph_def)
+    splitter.write(path)
+    chunks, _ = splitter.split()
+    if len(chunks) > 1:
+      logging.info("  %s.cpb written", path)
+    else:
+      raise RuntimeError(
+          "For some reason this graph was not chunked, so a .cpb file was not"
+          " produced. Raising an error since this should not be the case."
+      )
+
+
+def split_lots_nodes(path: str, export_files: Sequence[str]):
+  """GraphDef with lots of nodes."""
+  # The actual sizes in the generated graph has a slight deviation, but are
+  # between [90, 100] (tested in testMakeGraphDef with atol=5).
+  #    Expected Chunks (Max Size = 500)
+  #    -----------------------------
+  #       Chunk #: Contents
+  #    -----------------------------
+  #       0: GraphDef  # (nodes [0:5])
+  #    -----------------------------
+  #       1: GraphDef  # (nodes [5:10])
+  #    -----------------------------
+  #       2: GraphDef  # (nodes [10:15])
+  #    -----------------------------
+  #       3: ChunkedMessage
+  #    -----------------------------
+  graph_def = test_util.make_graph_def_with_constant_nodes(LOTS_NODES_SIZES)
+  _split_and_write(path, graph_def, 500, export_files)
+
+
+def split_large_nodes(path: str, export_files: Sequence[str]):
+  """GraphDef with large nodes."""
+  # Large nodes are greedily split from the original proto if they are
+  # larger than max_size / 3.
+  # This should create 6 chunks:
+  #   [parent GraphDef, node[1], node[2], node[3], node[5], ChunkedMessage]
+  graph_def = test_util.make_graph_def_with_constant_nodes(LARGE_NODES_SIZES)
+  _split_and_write(path, graph_def, 200, export_files)
+
+
+def split_large_constant(path: str, export_files: Sequence[str]):
+  """GraphDef with large constant nodes."""
+  #    Expected Chunks (Max Size = 500)
+  #    -----------------------------
+  #       Chunk #: Contents
+  #    -----------------------------
+  #       0: GraphDef
+  #    -----------------------------
+  #       1: GraphDef.nodes[2].attr["value"].tensor.tensor_content
+  #    -----------------------------
+  #       2: GraphDef.nodes[4].attr["value"].tensor.tensor_content
+  #    -----------------------------
+  graph_def = test_util.make_graph_def_with_constant_nodes(LARGE_CONSTANT_SIZES)
+  _split_and_write(path, graph_def, 500, export_files)
+
+
+def function_lots_of_nodes(path: str, export_files: Sequence[str]):
+  """Generates a proto of GraphDef with a FunctionDef that have many nodes."""
+  graph_def = test_util.make_graph_def_with_constant_nodes(
+      [], fn=LOTS_NODES_SIZES
+  )
+  _split_and_write(path, graph_def, 500, export_files)
+
+
+def function_large_nodes(path: str, export_files: Sequence[str]):
+  graph_def = test_util.make_graph_def_with_constant_nodes(
+      [], fn=LARGE_NODES_SIZES
+  )
+  _split_and_write(path, graph_def, 200, export_files)
+
+
+def graph_def_and_function(path: str, export_files: Sequence[str]):
+  graph_def = test_util.make_graph_def_with_constant_nodes(
+      [50, 50, 50, 50, 50, 50], fn1=[50, 50, 50], fn2=[50], fn3=[50], fn4=[50]
+  )
+  _split_and_write(path, graph_def, 200, export_files)
+
+
+VALID_GRAPH_TYPES = {
+    "split-lots-nodes": split_lots_nodes,
+    "split-large-nodes": split_large_nodes,
+    "split-large-constant": split_large_constant,
+    "function-lots-of-nodes": function_lots_of_nodes,
+    "function-large-nodes": function_large_nodes,
+    "graph-def-and-function": graph_def_and_function,
+}
+ALL_GRAPH_TYPES = ", ".join(VALID_GRAPH_TYPES.keys())
+
+SPLITTER_TESTDATA_PATH = flags.DEFINE_string(
+    "path", None, help="Path to testdata directory."
+)
+GRAPH_TYPES = flags.DEFINE_multi_string(
+    "graph_type",
+    "all",
+    help=f"Type(s) of graph to export. Valid types: all, {ALL_GRAPH_TYPES}",
+)
+EXPORT_FILES = flags.DEFINE_multi_string(
+    "export",
+    "all",
+    help="List of files to export. Valid options: all, pb, pbtxt, cpb",
+)
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+
+  if "all" in EXPORT_FILES.value:
+    export_files = ["pb", "pbtxt", "cpb"]
+  else:
+    export_files = EXPORT_FILES.value
+
+  if "all" in GRAPH_TYPES.value:
+    graph_types = VALID_GRAPH_TYPES.keys()
+  else:
+    graph_types = GRAPH_TYPES.value
+
+  for v in graph_types:
+    if v not in VALID_GRAPH_TYPES:
+      raise ValueError(
+          f"Invalid flag passed to `graph_type`: {v}\nValid graph types:"
+          f" {ALL_GRAPH_TYPES}"
+      )
+
+    logging.info("Generating graph %s", v)
+    f = VALID_GRAPH_TYPES[v]
+    f(os.path.join(SPLITTER_TESTDATA_PATH.value, v), export_files)
+
+
+if __name__ == "__main__":
+  app.run(main)

tensorflow/tools/proto_splitter/testdata/split_saved_model_gen.py

@@ -0,0 +1,145 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+r"""Generates SavedModel test data for Merger.
+
+Constructs chunked proto test data containing a SavedModel.
+
+Example command:
+
+bazel run tensorflow/tools/proto_splitter/testdata:split_saved_model_gen -- \
+  --path /tmp \
+  --saved_model_type=split-standard \
+  --export=pb,cpb
+"""
+
+from collections.abc import Sequence
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.python.lib.io import file_io
+from tensorflow.tools.proto_splitter import constants
+from tensorflow.tools.proto_splitter.python import saved_model as split_saved_model
+from tensorflow.tools.proto_splitter.python import test_util
+
+
+STANDARD_SIZES = [100, 100, 1000, 100, 1000, 500, 100, 100, 100]
+
+
+def _split_and_write(
+    path: str,
+    saved_model: saved_model_pb2.SavedModel,
+    max_size: int,
+    export_files: Sequence[str],
+):
+  """Writes the .pb, .pbtxt and .cpb files for a SavedModel."""
+  constants.debug_set_max_size(max_size)
+
+  if "pbtxt" in export_files:
+    output_path = f"{path}.pbtxt"
+    file_io.write_string_to_file(output_path, str(saved_model))
+    logging.info("  %s written", output_path)
+  if "pb" in export_files:
+    output_path = f"{path}.pb"
+    file_io.write_string_to_file(output_path, saved_model.SerializeToString())
+    logging.info("  %s written", output_path)
+  if "cpb" in export_files:
+    splitter = split_saved_model.SavedModelSplitter(saved_model)
+    splitter.write(path)
+    chunks, _ = splitter.split()
+    if len(chunks) > 1:
+      logging.info("  %s.cpb written", path)
+    else:
+      raise RuntimeError(
+          "For some reason this graph was not chunked, so a .cpb file was not"
+          " produced. Raising an error since this should not be the case."
+      )
+
+
+def split_standard(path: str, export_files: Sequence[str]):
+  """Splits a standard SavedModel."""
+  fn1 = [100, 100, 100]
+  fn2 = [100, 500]
+  fn3 = [100]
+  fn4 = [100, 100]
+
+  max_size = 500
+  constants.debug_set_max_size(max_size)
+
+  graph_def = test_util.make_graph_def_with_constant_nodes(
+      STANDARD_SIZES, fn1=fn1, fn2=fn2, fn3=fn3, fn4=fn4
+  )
+  proto = saved_model_pb2.SavedModel()
+  proto.meta_graphs.add().graph_def.CopyFrom(graph_def)
+
+  _split_and_write(path, proto, max_size, export_files)
+
+
+VALID_SAVED_MODEL_TYPES = {
+    "split-standard": split_standard,
+}
+ALL_SAVED_MODEL_TYPES = ", ".join(VALID_SAVED_MODEL_TYPES.keys())
+
+SPLITTER_TESTDATA_PATH = flags.DEFINE_string(
+    "path", None, help="Path to testdata directory."
+)
+SAVED_MODEL_TYPES = flags.DEFINE_multi_string(
+    "saved_model_type",
+    "all",
+    help=(
+        "Type(s) of saved model to export. Valid types: all, "
+        f"{ALL_SAVED_MODEL_TYPES}"
+    ),
+)
+EXPORT_FILES = flags.DEFINE_multi_string(
+    "export",
+    "all",
+    help="List of files to export. Valid options: all, pb, pbtxt, cpb",
+)
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+
+  if "all" in EXPORT_FILES.value:
+    export_files = ["pb", "pbtxt", "cpb"]
+  else:
+    export_files = EXPORT_FILES.value
+
+  if "all" in SAVED_MODEL_TYPES.value:
+    saved_model_types = VALID_SAVED_MODEL_TYPES.keys()
+  else:
+    saved_model_types = SAVED_MODEL_TYPES.value
+
+  for v in saved_model_types:
+    if v not in VALID_SAVED_MODEL_TYPES:
+      raise ValueError(
+          "Invalid flag passed to `saved_model_type`: "
+          f"{v}\nValid saved model types:"
+          f" {ALL_SAVED_MODEL_TYPES}"
+      )
+
+    logging.info("Generating saved model %s", v)
+    f = VALID_SAVED_MODEL_TYPES[v]
+    f(os.path.join(SPLITTER_TESTDATA_PATH.value, v), export_files)
+
+
+if __name__ == "__main__":
+  app.run(main)

tensorflow/tools/proto_splitter/testdata/test_message.proto

@@ -0,0 +1,42 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.proto_splitter_testdata;
+
+message RepeatedString {
+  repeated string strings = 1;
+}
+
+message RepeatedRepeatedString {
+  int32 filler_field = 1;
+  repeated RepeatedString rs = 2;
+}
+
+message ManyFields {
+  ManyFields field_one = 1;
+  repeated ManyFields repeated_field = 2;
+  string string_field = 3;
+  repeated string repeated_string_field = 4;
+  map<uint32, string> map_field_uint32 = 5;
+  map<int64, string> map_field_int64 = 6;
+  map<bool, ManyFields> nested_map_bool = 7;
+}
+
+message StringNode {
+  string val = 1;
+  repeated StringNode child_nodes = 2;
+}

tensorflow/tools/proto_splitter/util.py

@@ -0,0 +1,203 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Utilities for Proto Splitter modules."""
+
+from collections.abc import Sequence
+from typing import Any, Optional, Union
+
+from google.protobuf import descriptor
+from google.protobuf import message
+from tensorflow.tools.proto_splitter import chunk_pb2
+
+_BYTE_UNITS = [(1, "B"), (1 << 10, "KiB"), (1 << 20, "MiB"), (1 << 30, "GiB")]
+
+
+def format_bytes(b: int) -> str:
+  """Formats bytes into a human-readable string."""
+  for i in range(1, len(_BYTE_UNITS)):
+    if b < _BYTE_UNITS[i][0]:
+      n = f"{b / _BYTE_UNITS[i-1][0]:.2f}"
+      units = _BYTE_UNITS[i - 1][1]
+      break
+  else:
+    n = f"{b / _BYTE_UNITS[-1][0]:.2f}"
+    units = _BYTE_UNITS[-1][1]
+  n = n.rstrip("0").rstrip(".")
+  return f"{n}{units}"
+
+
+FieldTypes = Union[str, int, bool, Sequence[Union[str, int, bool]]]
+
+
+def get_field(
+    proto: message.Message, fields: FieldTypes
+) -> tuple[Any, Optional[descriptor.FieldDescriptor]]:
+  """Returns the field and field descriptor from the proto.
+
+  Args:
+    proto: Parent proto of any message type.
+    fields: List of string/int/map key fields, e.g. ["nodes", "attr", "value"]
+      can represent `proto.nodes.attr["value"]`.
+
+  Returns:
+    Tuple of (
+      Field in the proto or `None` if none are found,
+      Field descriptor
+    )
+  """
+  field_proto = proto
+  field_desc = None
+  for field_proto, field_desc, _, _ in _walk_fields(proto, fields):
+    pass
+  return field_proto, field_desc
+
+
+def get_field_tag(
+    proto: message.Message, fields: FieldTypes
+) -> Sequence[chunk_pb2.FieldIndex]:
+  """Generates FieldIndex proto for a nested field within a proto.
+
+  Args:
+    proto: Parent proto of any message type.
+    fields: List of string/int/map key fields, e.g. ["nodes", "attr", "value"]
+      can represent `proto.nodes.attr["value"]`.
+
+  Returns:
+    A list of FieldIndex protos with the same length as `fields`.
+  """
+  field_tags = []
+  for _, field_desc, map_key, list_index in _walk_fields(proto, fields):
+    field_tags.append(chunk_pb2.FieldIndex(field=field_desc.number))
+    if map_key is not None:
+      key_type = field_desc.message_type.fields_by_name["key"].type
+      field_tags.append(
+          chunk_pb2.FieldIndex(map_key=_map_key_proto(key_type, map_key))
+      )
+    elif list_index is not None:
+      field_tags.append(chunk_pb2.FieldIndex(index=list_index))
+  return field_tags
+
+
+def _walk_fields(proto: message.Message, fields: FieldTypes):
+  """Yields fields in a proto.
+
+  Args:
+    proto: Parent proto of any message type.
+    fields: List of string/int/map key fields, e.g. ["nodes", "attr", "value"]
+      can represent `proto.nodes.attr["value"]`.
+
+  Yields:
+    Tuple of (
+      Field in the proto or `None` if none are found,
+      Field descriptor,
+      Key into this map field (or None),
+      Index into this repeated field (or None))
+  """
+  if not isinstance(fields, list):
+    fields = [fields]
+
+  field_proto = proto
+  parent_desc = proto.DESCRIPTOR
+  i = 0
+  while i < len(fields):
+    field = fields[i]
+    field_desc = None
+    map_key = None
+    index = None
+
+    if parent_desc is None:
+      raise ValueError(
+          f"Unable to find fields: {fields} in proto of type {type(proto)}."
+      )
+
+    if isinstance(field, int):
+      try:
+        field_desc = parent_desc.fields_by_number[field]
+      except KeyError:
+        raise KeyError(  # pylint:disable=raise-missing-from
+            f"Unable to find field number {field} in {parent_desc.full_name}. "
+            f"Valid field numbers: {parent_desc.fields_by_number.keys()}"
+        )
+    elif isinstance(field, str):
+      try:
+        field_desc = parent_desc.fields_by_name[field]
+      except KeyError:
+        raise KeyError(  # pylint:disable=raise-missing-from
+            f"Unable to find field '{field}' in {parent_desc.full_name}. "
+            f"Valid field names: {parent_desc.fields_by_name.keys()}"
+        )
+    else:  # bool (only expected as map key)
+      raise TypeError("Unexpected bool found in field list.")
+
+    i += 1
+    parent_desc = field_desc.message_type
+    if field_proto is not None:
+      field_proto = getattr(field_proto, field_desc.name)
+
+    # Handle special fields types (map key and list index).
+    if _is_map(parent_desc) and i < len(fields):
+      # Next field is the map key.
+      map_key = fields[i]
+
+      try:
+        field_proto = field_proto[map_key] if field_proto is not None else None
+      except KeyError:
+        field_proto = None
+      i += 1
+
+      if i < len(fields):
+        # The next field must be from the Value Message.
+        value_desc = parent_desc.fields_by_name["value"]
+        assert value_desc.message_type is not None
+        parent_desc = value_desc.message_type
+
+    elif is_repeated(field_desc) and i < len(fields):
+      # The next field is the index within the list.
+      index = fields[i]
+      try:
+        field_proto = field_proto[index] if field_proto is not None else None
+      except IndexError:
+        field_proto = None
+      i += 1
+
+    yield field_proto, field_desc, map_key, index
+
+
+def _is_map(desc: descriptor.Descriptor) -> bool:
+  return desc.GetOptions().map_entry if desc is not None else False
+
+
+def is_repeated(field_desc: descriptor.FieldDescriptor) -> bool:
+  # copybara:uncomment_begin(google-only)
+  # return field_desc.is_repeated
+  # copybara:uncomment_end_and_comment_begin
+  return field_desc.label == descriptor.FieldDescriptor.LABEL_REPEATED
+  # copybara:comment_end
+
+
+_FIELD_DESC = descriptor.FieldDescriptor
+_MAP_KEY = {
+    _FIELD_DESC.TYPE_STRING: lambda key: chunk_pb2.FieldIndex.MapKey(s=key),
+    _FIELD_DESC.TYPE_BOOL: lambda key: chunk_pb2.FieldIndex.MapKey(boolean=key),
+    _FIELD_DESC.TYPE_UINT32: lambda key: chunk_pb2.FieldIndex.MapKey(ui32=key),
+    _FIELD_DESC.TYPE_UINT64: lambda key: chunk_pb2.FieldIndex.MapKey(ui64=key),
+    _FIELD_DESC.TYPE_INT32: lambda key: chunk_pb2.FieldIndex.MapKey(i32=key),
+    _FIELD_DESC.TYPE_INT64: lambda key: chunk_pb2.FieldIndex.MapKey(i64=key),
+}
+
+
+def _map_key_proto(key_type, key):
+  """Returns MapKey proto for a key of key_type."""
+  return _MAP_KEY[key_type](key)

tensorflow/tools/proto_splitter/util_test.py

@@ -0,0 +1,135 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Tests util functions."""
+
+from typing import Iterable
+
+from tensorflow.python.platform import test
+from tensorflow.tools.proto_splitter import util
+from tensorflow.tools.proto_splitter.testdata import test_message_pb2
+
+
+class UtilTest(test.TestCase):
+
+  def test_format_bytes(self):
+    self.assertEqual(util.format_bytes(1024), "1KiB")
+    self.assertEqual(util.format_bytes(5632), "5.5KiB")
+    self.assertEqual(util.format_bytes(53432), "52.18KiB")
+
+    self.assertEqual(util.format_bytes(76493281), "72.95MiB")
+    self.assertEqual(util.format_bytes(5.977e7), "57MiB")
+    self.assertEqual(util.format_bytes(1.074e9), "1GiB")
+    self.assertEqual(util.format_bytes(16493342281), "15.36GiB")
+
+  def test_get_field_tag(self):
+    proto = test_message_pb2.ManyFields()
+
+    # proto.field_one.repeated_field[15].map_field_uint32[10]
+    ret = util.get_field_tag(
+        proto, ["field_one", "repeated_field", 15, "map_field_uint32", 10]
+    )
+    self.assertLen(ret, 5)
+    self.assertEqual(1, ret[0].field)
+    self.assertEqual(2, ret[1].field)
+    self.assertEqual(15, ret[2].index)
+    self.assertEqual(5, ret[3].field)
+    self.assertEqual(10, ret[4].map_key.ui32)
+    self.assertFalse(ret[4].map_key.HasField("i32"))
+
+    # proto.nested_map_bool[False].map_field_int64
+    ret = util.get_field_tag(
+        proto, ["nested_map_bool", False, "map_field_int64"]
+    )
+
+    self.assertLen(ret, 3)
+    self.assertEqual(7, ret[0].field)
+    self.assertEqual(False, ret[1].map_key.boolean)
+    self.assertEqual(6, ret[2].field)
+
+    # proto.repeated_field[55].nested_map_bool[True].string_field
+    ret = util.get_field_tag(proto, [2, 55, 7, True, 3])
+    self.assertLen(ret, 5)
+    self.assertEqual(2, ret[0].field)
+    self.assertEqual(55, ret[1].index)
+    self.assertEqual(7, ret[2].field)
+    self.assertEqual(True, ret[3].map_key.boolean)
+    self.assertEqual(3, ret[4].field)
+
+  def test_get_field_tag_invalid(self):
+    proto = test_message_pb2.ManyFields()
+    with self.assertRaisesRegex(KeyError, "Unable to find field 'not_a_field'"):
+      util.get_field_tag(proto, ["field_one", "not_a_field"])
+    with self.assertRaisesRegex(KeyError, "Unable to find field number 10000"):
+      util.get_field_tag(proto, [1, 10000])
+    with self.assertRaisesRegex(ValueError, "Unable to find fields.*in proto"):
+      util.get_field_tag(proto, ["string_field", 1])
+
+  def test_get_field_and_desc(self):
+    proto = test_message_pb2.ManyFields(
+        field_one=test_message_pb2.ManyFields(
+            repeated_field=[
+                test_message_pb2.ManyFields(),
+                test_message_pb2.ManyFields(
+                    string_field="inner_inner_string",
+                    map_field_uint32={
+                        324: "map_value_324",
+                        543: "map_value_543",
+                    },
+                ),
+            ]
+        ),
+        map_field_int64={
+            -1345: "map_value_-1345",
+        },
+        nested_map_bool={
+            True: test_message_pb2.ManyFields(string_field="string_true"),
+            False: test_message_pb2.ManyFields(string_field="string_false"),
+        },
+    )
+
+    field, field_desc = util.get_field(proto, [])
+    self.assertIs(proto, field)
+    self.assertIsNone(field_desc)
+
+    field, field_desc = util.get_field(proto, ["field_one", "repeated_field"])
+    self.assertIsInstance(field, Iterable)
+    self.assertLen(field, 2)
+    self.assertEqual("repeated_field", field_desc.name)
+    self.assertEqual(2, field_desc.number)
+    self.assertProtoEquals(proto.field_one.repeated_field, field)
+
+    field, field_desc = util.get_field(proto, ["field_one", 2, 1])
+    self.assertIsInstance(field, test_message_pb2.ManyFields)
+    self.assertEqual("repeated_field", field_desc.name)
+    self.assertEqual(2, field_desc.number)
+    self.assertProtoEquals(proto.field_one.repeated_field[1], field)
+
+    field, _ = util.get_field(proto, ["field_one", 2, 1, "string_field"])
+    self.assertEqual("inner_inner_string", field)
+
+    field, _ = util.get_field(
+        proto, ["field_one", 2, 1, "map_field_uint32", 324]
+    )
+    self.assertEqual("map_value_324", field)
+
+    field, _ = util.get_field(proto, ["nested_map_bool", False, "string_field"])
+    self.assertEqual("string_false", field)
+
+    field, _ = util.get_field(proto, ["nested_map_bool", True, "string_field"])
+    self.assertEqual("string_true", field)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/tools/proto_splitter/version.py

@@ -0,0 +1,32 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Version constants."""
+
+# Version of the Join implementation.
+# This should be incremented almost whenever the joining implementation is
+# updated. The only time this number should not be incremented is if the change
+# is extremely trivial.
+_JOIN_VERSION = 0
+
+# Bad/Buggy join versions.
+_BAD_VERSIONS = ()
+
+
+def get_current_join_version() -> int:
+  return _JOIN_VERSION
+
+
+def get_bad_versions() -> tuple[int, ...]:
+  return _BAD_VERSIONS

tensorflow/tools/proto_splitter/versions.proto

@@ -0,0 +1,32 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.proto_splitter;
+
+option cc_enable_arenas = true;
+
+// Version information for Chunked protos.
+message VersionDef {
+  // Information about the Splitter used to split and write this data.
+  int32 splitter_version = 1;
+
+  // The version of the Join implemention required to consume this data.
+  int32 join_version = 2;
+
+  // Specific consumer versions which are disallowed (e.g. due to bugs).
+  repeated int32 bad_consumers = 3;
+}