Skip to content

Stop statically linking TensorFlow into every OSS TF.Text op library. - #1572

Merged
copybara-service[bot] merged 1 commit into
masterfrom
test_1002376850
Oct 23, 2026
Merged

copybara-service[bot] merged 1 commit into
masterfrom
test_1002376850

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Oct 19, 2026 •

Copy link
Copy Markdown
Contributor

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 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.

@copybara-service
copybara-service Bot force-pushed the test_1002376850 branch 9 times, most recently from 7364b60 to c24fae7 Compare October 23, 2026 14:15
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
copybara-service Bot merged commit 738157d into master Oct 23, 2026
4 checks passed
@copybara-service
copybara-service Bot deleted the test_1002376850 branch October 23, 2026 14:32

This branch was successfully deployed

1 active deployment
github-pages — 738157d5 Deployed Oct 23, 2026 by github-pages[bot] via deploy #338
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant