chore: import upstream snapshot with attribution

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()