From 738157d56ec991e0bba37b9457f258f7a639d697 Mon Sep 17 00:00:00 2001 From: "TF.Text Team" Date: Fri, 23 Oct 2026 07:31:40 -0700 Subject: [PATCH] Stop statically linking TensorFlow into every OSS TF.Text op library. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since TF 2.20 the TensorFlow build enables pywrap rules, so `if_static()` resolves to full implementations, and TF proto targets are no longer header-only. Because every OSS `py_tf_text_library` depended on `@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim`, each `_.so` statically linked its own copy of TF protos, the protobuf runtime and absl (5–9 MB per op library). That inflated the wheel about 8x and put duplicate TensorFlow, protobuf and absl state into the process. Compile only the three shim sources against the TensorFlow wheel headers and link them dynamically against `libtensorflow_framework`, like any other custom op. A small patch exports the shim sources from the TensorFlow source tree and keeps the `AttrSlice::Find` lookup in `TfInitContext::GetAttr`, which is required because `Map::at()` on a `NodeDef` owned by `libtensorflow_framework` fails across the shared-library boundary. With that, `tensorflow.xla_if_static.bzl.patch` no longer affects the op libraries; it only added TF protos to `tflite_op_shim`, which `libpybinds_library_common.so` links as a whole archive. Remove the patch together with the unused TensorFlow core includes that still required it: `whitespace_tokenizer_kernel_template.h` no longer includes `tensor.h`, `sentence_fragmenter_v2.{h,cc}` no longer include `status.h` (the struct uses `uint32_t` and the test `std::string` instead of the TF aliases), and `fast_wordpiece_tokenizer_model_builder.cc` includes `absl/log/log.h` directly instead of getting `LOG` transitively. PiperOrigin-RevId: 1005586601 --- WORKSPACE | 2 +- .../fast_wordpiece_tokenizer_model_builder.cc | 1 + .../core/kernels/sentence_fragmenter_v2.cc | 1 - .../core/kernels/sentence_fragmenter_v2.h | 4 +- .../kernels/sentence_fragmenter_v2_test.cc | 42 ++++++------ .../whitespace_tokenizer_kernel_template.h | 1 - tensorflow_text/tftext.bzl | 27 +++++++- .../tensorflow/tensorflow.shim_sources.patch | 33 ++++++++++ .../tensorflow.xla_if_static.bzl.patch | 64 ------------------- 9 files changed, 84 insertions(+), 91 deletions(-) create mode 100644 third_party/tensorflow/tensorflow.shim_sources.patch delete mode 100644 third_party/tensorflow/tensorflow.xla_if_static.bzl.patch diff --git a/WORKSPACE b/WORKSPACE index b4653a2bf..08b0196bc 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -72,7 +72,7 @@ http_archive( name = "org_tensorflow", patch_args = ["-p1"], patches = [ - "//third_party/tensorflow:tensorflow.xla_if_static.bzl.patch", + "//third_party/tensorflow:tensorflow.shim_sources.patch", ], sha256 = "b7fc9646df01c14968e1546ad430a72f6ecc26867c085266f71f43968a6d5f5b", strip_prefix = "tensorflow-2.22.0-rc0", diff --git a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc index 5aed0bb42..5e44e0fc4 100644 --- a/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc +++ b/tensorflow_text/core/kernels/fast_wordpiece_tokenizer_model_builder.cc @@ -22,6 +22,7 @@ #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" diff --git a/tensorflow_text/core/kernels/sentence_fragmenter_v2.cc b/tensorflow_text/core/kernels/sentence_fragmenter_v2.cc index d917106c2..983324c71 100644 --- a/tensorflow_text/core/kernels/sentence_fragmenter_v2.cc +++ b/tensorflow_text/core/kernels/sentence_fragmenter_v2.cc @@ -21,7 +21,6 @@ #include "absl/strings/string_view.h" #include "icu4c/source/common/unicode/uchar.h" #include "icu4c/source/common/unicode/utf8.h" -#include "tensorflow/core/lib/core/status.h" namespace tensorflow { namespace text { diff --git a/tensorflow_text/core/kernels/sentence_fragmenter_v2.h b/tensorflow_text/core/kernels/sentence_fragmenter_v2.h index 6c06867eb..59e9b5603 100644 --- a/tensorflow_text/core/kernels/sentence_fragmenter_v2.h +++ b/tensorflow_text/core/kernels/sentence_fragmenter_v2.h @@ -34,12 +34,12 @@ #ifndef TENSORFLOW_TEXT_CORE_KERNELS_SENTENCE_FRAGMENTER_V2_H_ #define TENSORFLOW_TEXT_CORE_KERNELS_SENTENCE_FRAGMENTER_V2_H_ +#include #include #include "absl/status/status.h" #include "absl/strings/string_view.h" #include "icu4c/source/common/unicode/utypes.h" -#include "tensorflow/core/lib/core/status.h" namespace tensorflow { namespace text { @@ -140,7 +140,7 @@ struct SentenceFragment { HAS_SENTENTIAL_CLOSE_PAREN = 0x0008, // e.g.: (Mushrooms are fungi!) }; // A mask of the above listed properties. - uint32 properties = 0; + uint32_t properties = 0; int terminal_punc_token = -1; }; diff --git a/tensorflow_text/core/kernels/sentence_fragmenter_v2_test.cc b/tensorflow_text/core/kernels/sentence_fragmenter_v2_test.cc index 87cd49265..e7a9ad030 100644 --- a/tensorflow_text/core/kernels/sentence_fragmenter_v2_test.cc +++ b/tensorflow_text/core/kernels/sentence_fragmenter_v2_test.cc @@ -736,7 +736,7 @@ TEST_F(IsEmoticonTest, IsEmoticon) { TEST(SentenceFragmenterTest, Basic) { // 1 // 012345678901234 - string test_input = "Hello. Foo bar!"; + std::string test_input = "Hello. Foo bar!"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -749,7 +749,7 @@ TEST(SentenceFragmenterTest, Basic) { TEST(SentenceFragmenterTest, BasicEllipsis) { // 1 // 012345678901234 - string test_input = "Hello...foo bar"; + std::string test_input = "Hello...foo bar"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -763,7 +763,7 @@ TEST(SentenceFragmenterTest, BasicEllipsis) { TEST(SentenceFragmenterTest, Parentheses) { // 1 2 // 012345678901234567890123456789 - string test_input = "Hello (who are you...) foo bar"; + std::string test_input = "Hello (who are you...) foo bar"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -776,7 +776,7 @@ TEST(SentenceFragmenterTest, Parentheses) { TEST(SentenceFragmenterTest, MidFragmentParentheses) { // 1 2 // 012345678901234567890123456789 - string test_input = "Hello (who are you) world? Foo bar"; + std::string test_input = "Hello (who are you) world? Foo bar"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -789,7 +789,7 @@ TEST(SentenceFragmenterTest, MidFragmentParentheses) { TEST(SentenceFragmenterTest, PunctuationAfterParentheses) { // 1 2 // 01234567890123456789012345678 - string test_input = "Hello (who are you)? Foo bar!"; + std::string test_input = "Hello (who are you)? Foo bar!"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -802,7 +802,7 @@ TEST(SentenceFragmenterTest, PunctuationAfterParentheses) { TEST(SentenceFragmenterTest, ManyFinalPunctuations) { // 1 2 // 0123456789012345678901234 - string test_input = "Hello!!!!! Who are you??"; + std::string test_input = "Hello!!!!! Who are you??"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -815,7 +815,7 @@ TEST(SentenceFragmenterTest, ManyFinalPunctuations) { TEST(SentenceFragmenterTest, NewLine) { // 1 2 3 // 012345678901234567890 1 23456 7 89012 3 45678 - string test_input = "Who let the dogs out?\r\nWho?\r\nWho?\r\nWho?"; + std::string test_input = "Who let the dogs out?\r\nWho?\r\nWho?\r\nWho?"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -832,7 +832,7 @@ TEST(SentenceFragmenterTest, NewLine) { TEST(SentenceFragmenterTest, WhiteSpaceInPunctuation) { // 1 2 // 0123456789012345678901234 - string test_input = "Hello?? !!! Who are you??"; + std::string test_input = "Hello?? !!! Who are you??"; SentenceFragmenterV2 fragmenter(test_input); std::vector fragments; EXPECT_TRUE(fragmenter.FindFragments(&fragments).ok()); @@ -850,7 +850,7 @@ TEST(FragmentBoundaryMatchTest, NoStateChange) { FragmentBoundaryMatch f; // || // 012345678901234 - string test_input = "Hello...foo bar"; + std::string test_input = "Hello...foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_FALSE(f.GotTerminalPunc()); @@ -864,7 +864,7 @@ TEST(FragmentBoundaryMatchTest, BasicEllipsis) { FragmentBoundaryMatch f; // | | // 0123456789 - string test_input = "...foo bar"; + std::string test_input = "...foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -878,7 +878,7 @@ TEST(FragmentBoundaryMatchTest, BasicPeriod) { FragmentBoundaryMatch f; // || // 0123456789 - string test_input = ". Foo bar"; + std::string test_input = ". Foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -892,7 +892,7 @@ TEST(FragmentBoundaryMatchTest, BasicAcronym) { FragmentBoundaryMatch f; // | | // 0123456789 - string test_input = "A.B. xyz"; + std::string test_input = "A.B. xyz"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -906,7 +906,7 @@ TEST(FragmentBoundaryMatchTest, LongerAcronym) { FragmentBoundaryMatch f; // | | // 0123456789 - string test_input = "I.B.M. yo"; + std::string test_input = "I.B.M. yo"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -920,7 +920,7 @@ TEST(FragmentBoundaryMatchTest, Emoticon) { FragmentBoundaryMatch f; // | | // 0123456789012 - string test_input = ">:-( hello..."; + std::string test_input = ">:-( hello..."; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -934,7 +934,7 @@ TEST(FragmentBoundaryMatchTest, ParensWithEllipsis) { FragmentBoundaryMatch f; // || // 0123456789012345 - string test_input = ".foo...) foo bar"; + std::string test_input = ".foo...) foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -948,7 +948,7 @@ TEST(FragmentBoundaryMatchTest, ClosingParenWithEllipsis) { FragmentBoundaryMatch f; // | | // 012345678901 - string test_input = "...) foo bar"; + std::string test_input = "...) foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -962,7 +962,7 @@ TEST(FragmentBoundaryMatchTest, BeginAndEndParenWithEllipsis) { FragmentBoundaryMatch f; // || // 0123456789012 - string test_input = "(...) foo bar"; + std::string test_input = "(...) foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_FALSE(f.GotTerminalPunc()); @@ -986,7 +986,7 @@ TEST(FragmentBoundaryMatchTest, AcronymInSentence) { FragmentBoundaryMatch f; // | | // 0123456789012 - string test_input = "U.S. don't be surprised."; + std::string test_input = "U.S. don't be surprised."; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); @@ -1000,7 +1000,7 @@ TEST(FragmentBoundaryMatchTest, HelloWithEllipsis) { FragmentBoundaryMatch f; // || // 01234567890 - string test_input = "o...foo bar"; + std::string test_input = "o...foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_FALSE(f.GotTerminalPunc()); @@ -1024,7 +1024,7 @@ TEST(FragmentBoundaryMatchTest, ThreeStatesWithClosigParen) { FragmentBoundaryMatch f; // || // 0123456789012 - string test_input = "w...) foo bar"; + std::string test_input = "w...) foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_FALSE(f.GotTerminalPunc()); @@ -1068,7 +1068,7 @@ TEST(FragmentBoundaryMatchTest, NoTransition) { FragmentBoundaryMatch f; // | | // 0123456789012 - string test_input = "...foo bar"; + std::string test_input = "...foo bar"; int index = 0; EXPECT_TRUE(f.Advance(index, test_input)); EXPECT_TRUE(f.GotTerminalPunc()); diff --git a/tensorflow_text/core/kernels/whitespace_tokenizer_kernel_template.h b/tensorflow_text/core/kernels/whitespace_tokenizer_kernel_template.h index e05a914bb..214072b64 100644 --- a/tensorflow_text/core/kernels/whitespace_tokenizer_kernel_template.h +++ b/tensorflow_text/core/kernels/whitespace_tokenizer_kernel_template.h @@ -20,7 +20,6 @@ #include "absl/status/status.h" #include "absl/status/statusor.h" -#include "tensorflow/core/framework/tensor.h" #include "tensorflow/lite/kernels/shim/op_kernel.h" #include "tensorflow/lite/kernels/shim/shape.h" #include "tensorflow/lite/kernels/shim/status_macros.h" diff --git a/tensorflow_text/tftext.bzl b/tensorflow_text/tftext.bzl index 1c18c6263..44ab3df73 100644 --- a/tensorflow_text/tftext.bzl +++ b/tensorflow_text/tftext.bzl @@ -9,6 +9,31 @@ load( _pywrap_library = "pywrap_library", ) +def _local_tf_op_shim(): + """Returns a package-local `tf_op_shim` built against the TF wheel headers. + + Only the three shim sources are compiled here. They are linked dynamically + against libtensorflow_framework.so (like any TF custom op), instead of + depending on @org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim, which + statically pulls TensorFlow implementation targets (protos, protobuf + runtime, lib_internal, ...) into every _.so. + """ + name = "tf_op_shim_oss" + if not native.existing_rule(name): + tf_cc_library( + name = name, + srcs = [ + "@org_tensorflow//tensorflow/lite/kernels/shim:shape.cc", + "@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim.cc", + "@org_tensorflow//tensorflow/lite/kernels/shim:tf_tensor_view.cc", + ], + deps = [ + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/types:span", + ], + ) + return ":" + name + def py_tf_text_library( name, srcs = [], @@ -55,7 +80,7 @@ def py_tf_text_library( }), alwayslink = 1, deps = cc_op_kernels + - ["@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim"] + + [_local_tf_op_shim()] + select({ "@org_tensorflow//tensorflow:mobile": [ "@org_tensorflow//tensorflow/core:portable_tensorflow_lib_lite", diff --git a/third_party/tensorflow/tensorflow.shim_sources.patch b/third_party/tensorflow/tensorflow.shim_sources.patch new file mode 100644 index 000000000..d32b28542 --- /dev/null +++ b/third_party/tensorflow/tensorflow.shim_sources.patch @@ -0,0 +1,33 @@ +diff --git a/tensorflow/lite/kernels/shim/BUILD b/tensorflow/lite/kernels/shim/BUILD +--- a/tensorflow/lite/kernels/shim/BUILD ++++ b/tensorflow/lite/kernels/shim/BUILD +@@ -13,6 +13,15 @@ + licenses = ["notice"], + ) + ++exports_files( ++ [ ++ "shape.cc", ++ "tf_op_shim.cc", ++ "tf_tensor_view.cc", ++ ], ++ visibility = ["//visibility:public"], ++) ++ + filegroup( + name = "tf_headers", + srcs = [ +diff --git a/tensorflow/lite/kernels/shim/tf_op_shim.cc b/tensorflow/lite/kernels/shim/tf_op_shim.cc +--- a/tensorflow/lite/kernels/shim/tf_op_shim.cc ++++ b/tensorflow/lite/kernels/shim/tf_op_shim.cc +@@ -79,8 +79,8 @@ absl::StatusOr TfInitContext::GetAttr( + return absl::InvalidArgumentError( + absl::StrCat("Non-existent attribute: ", attr_name, "\nop def:\n", + context_->def().DebugString())); +- const auto& attr_value = context_->def().attr().at(attr_name); +- return TfAttrValueToShimAttrValue(attr_value); ++ const auto* attr_value = ::tensorflow::AttrSlice(context_->def()).Find(attr_name); ++ return TfAttrValueToShimAttrValue(*attr_value); + } + + TfInvokeContext::TfInvokeContext(::tensorflow::OpKernelContext* context) diff --git a/third_party/tensorflow/tensorflow.xla_if_static.bzl.patch b/third_party/tensorflow/tensorflow.xla_if_static.bzl.patch deleted file mode 100644 index 5a5df740c..000000000 --- a/third_party/tensorflow/tensorflow.xla_if_static.bzl.patch +++ /dev/null @@ -1,64 +0,0 @@ -diff --git a/tensorflow/lite/kernels/shim/BUILD b/tensorflow/lite/kernels/shim/BUILD ---- a/tensorflow/lite/kernels/shim/BUILD -+++ b/tensorflow/lite/kernels/shim/BUILD -@@ -187,6 +187,9 @@ cc_library( - ":status_macros", - ":tensor_view", - ":tflite_tensor_view", -+ "//tensorflow/core:framework_headers_lib", -+ "//tensorflow/core:lib_headers_for_pybind", -+ "//tensorflow/core:protos_all_cc", - "//tensorflow/lite:mutable_op_resolver", - "//tensorflow/lite:string_util", - "//tensorflow/lite/core/c:common", -diff --git a/tensorflow/lite/kernels/shim/tf_op_shim.cc b/tensorflow/lite/kernels/shim/tf_op_shim.cc ---- a/tensorflow/lite/kernels/shim/tf_op_shim.cc -+++ b/tensorflow/lite/kernels/shim/tf_op_shim.cc -@@ -79,8 +79,8 @@ absl::StatusOr TfInitContext::GetAttr( - return absl::InvalidArgumentError( - absl::StrCat("Non-existent attribute: ", attr_name, "\nop def:\n", - context_->def().DebugString())); -- const auto& attr_value = context_->def().attr().at(attr_name); -- return TfAttrValueToShimAttrValue(attr_value); -+ const auto* attr_value = ::tensorflow::AttrSlice(context_->def()).Find(attr_name); -+ return TfAttrValueToShimAttrValue(*attr_value); - } - - TfInvokeContext::TfInvokeContext(::tensorflow::OpKernelContext* context) -diff --git a/third_party/xla/xla/tsl/platform/default/build_config.bzl b/third_party/xla/xla/tsl/platform/default/build_config.bzl ---- a/third_party/xla/xla/tsl/platform/default/build_config.bzl -+++ b/third_party/xla/xla/tsl/platform/default/build_config.bzl -@@ -390,6 +390,9 @@ def tf_additional_lib_deps(): - return [ - "@com_google_absl//absl/base:base", - "@com_google_absl//absl/container:inlined_vector", -+ "@com_google_absl//absl/status:status_macros", - "@com_google_absl//absl/types:span", -+ "@tsl//tsl/platform:protobuf", -+ "@xla//xla/tsl/concurrency:executor", - ] - - def tf_additional_core_deps(): -diff --git a/third_party/xla/xla/tsl/platform/default/build_config_root.bzl b/third_party/xla/xla/tsl/platform/default/build_config_root.bzl ---- a/third_party/xla/xla/tsl/platform/default/build_config_root.bzl -+++ b/third_party/xla/xla/tsl/platform/default/build_config_root.bzl -@@ -104,9 +104,6 @@ def if_google(google_value, oss_value = []): - # used for all framework_shared_object platforms including MacOS. - # buildifier: disable=function-docstring - def if_static(extra_deps, otherwise = [], macos = []): -- if use_pywrap_rules(): -- return extra_deps -- - ret = { - str(Label("//xla/tsl:framework_shared_object")): otherwise, - "//conditions:default": extra_deps, -@@ -117,9 +114,6 @@ def if_static(extra_deps, otherwise = [], macos = []): - - # TODO(b/356020232): remove completely after migration is done - def if_static_and_not_mobile(extra_deps, otherwise = []): -- if use_pywrap_rules(): -- return extra_deps -- - return select({ - str(Label("//xla/tsl:framework_shared_object")): otherwise, - str(Label("//xla/tsl:android")): otherwise,