Repository navigation
Stop statically linking TensorFlow into every OSS TF.Text op library. - #1572
Merged
Merged
Conversation
copybara-service
Bot
force-pushed
the
test_1002376850
branch
9 times, most recently
from
October 23, 2026 14:15
7364b60 to
c24fae7
Compare
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 `_<op>.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
copybara-service
Bot
force-pushed
the
test_1002376850
branch
from
October 23, 2026 14:32
c24fae7 to
738157d
Compare
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stop statically linking TensorFlow into every OSS TF.Text op library.
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 OSSpy_tf_text_librarydepended on@org_tensorflow//tensorflow/lite/kernels/shim:tf_op_shim, each_<op>.sostatically 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 theAttrSlice::Findlookup inTfInitContext::GetAttr, which is required becauseMap::at()on aNodeDefowned bylibtensorflow_frameworkfails across the shared-library boundary.With that,
tensorflow.xla_if_static.bzl.patchno longer affects the op libraries; it only added TF protos totflite_op_shim, whichlibpybinds_library_common.solinks as a whole archive. Remove the patch together with the unused TensorFlow core includes that still required it:whitespace_tokenizer_kernel_template.hno longer includestensor.h,sentence_fragmenter_v2.{h,cc}no longer includestatus.h(the struct usesuint32_tand the teststd::stringinstead of the TF aliases), andfast_wordpiece_tokenizer_model_builder.ccincludesabsl/log/log.hdirectly instead of gettingLOGtransitively.