part#1 Use pywrap rules for LiteRT repo to fix duplicate registration errors

PiperOrigin-RevId: 767324105

Author: ecalubaquib

tflite/python/BUILD

@@ -1,7 +1,8 @@
 load("@flatbuffers//:build_defs.bzl", "flatbuffer_py_library")
 load("@org_tensorflow//tensorflow:pytype.default.bzl", "pytype_strict_contrib_test", "pytype_strict_library")
 load("@org_tensorflow//tensorflow:strict.default.bzl", "py_strict_binary", "py_strict_library", "py_strict_test")
-load("@org_tensorflow//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("@org_tensorflow//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "pywrap_library")
+load("@org_tensorflow//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
 load("@rules_shell//shell:sh_test.bzl", "sh_test")
 load("//tflite:special_rules.bzl", "internal_visibility_allowlist")
 
@@ -474,7 +475,9 @@ pytype_strict_library(
         "@org_tensorflow//tensorflow/python/framework:tensor_shape",
         "@org_tensorflow//tensorflow/python/util:deprecation",
         "@org_tensorflow//tensorflow/python/util:tf_export",
-    ],
+    ] + if_pywrap(
+        if_true = [":pywrap_tflite"],
+    ),
 )
 
 py_strict_library(
@@ -515,7 +518,9 @@ pytype_strict_contrib_test(
         "@org_tensorflow//tensorflow/python/ops:math_ops",
         "@org_tensorflow//tensorflow/python/platform:client_testlib",
         "@org_tensorflow//third_party/py/numpy",
-    ],
+    ] + if_pywrap(
+        if_true = [":pywrap_tflite"],
+    ),
 )
 
 py_strict_library(
@@ -608,10 +613,16 @@ py_strict_library(
     ],
     visibility = ["//visibility:public"],
     deps = [
-        "//tflite/python/analyzer_wrapper:_pywrap_analyzer_wrapper",
         "@org_tensorflow//tensorflow/compiler/mlir/lite/python:wrap_converter",
         "@org_tensorflow//tensorflow/python/util:tf_export",
-    ],
+    ] + if_pywrap(
+        if_false = [
+            "//tflite/python/analyzer_wrapper:_pywrap_analyzer_wrapper",
+        ],
+        if_true = [
+            ":pywrap_tflite",
+        ],
+    ),
 )
 
 py_strict_test(
@@ -633,3 +644,27 @@ py_strict_test(
         "@org_tensorflow//tensorflow/python/trackable:autotrackable",
     ],
 )
+
+# Use pywrap_library to avoid duplicate registration of pybind11 modules.
+# A great example on how to use pywrap_library is
+# https://github.com/vam-google/symbol-locations/blob/main/pybind/BUILD
+# The following pywrap_library is used by LiteRT repo to avoid shared links provided
+# by Tensorflow under tensorflow/python:_pywrap_tensorflow
+# This isolate LiteRT's pybind11 dependencies. To use, add pybind deps under pywrap_tflite
+# and refer pywrap_tflite to any target that needsd to selected shared objects.
+pywrap_library(
+    name = "pywrap_tflite",
+    common_lib_versions = {
+        "framework": "2",
+    },
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tflite/python/analyzer_wrapper:_pywrap_analyzer_wrapper",
+        "//tflite/python/interpreter_wrapper:_pywrap_tensorflow_interpreter_wrapper",
+        "//tflite/python/metrics:_pywrap_tensorflow_lite_metrics_wrapper",
+        "//tflite/python/optimize:_pywrap_tensorflow_lite_calibration_wrapper",
+        "//tflite/testing:_pywrap_string_util",
+        "//tflite/tools/optimize/python:_pywrap_modify_model_interface",
+        "//tflite/tools/optimize/sparsity:format_converter_wrapper_pybind11",
+    ],
+)

tflite/python/metrics/BUILD

@@ -36,7 +36,7 @@ pybind_extension(
     pytype_srcs = [
         "_pywrap_tensorflow_lite_metrics_wrapper.pyi",
     ],
-    visibility = ["__subpackages__"],
+    visibility = ["//visibility:public"],
     deps = [
         ":metrics_wrapper_lib",
         "@com_google_protobuf//:protobuf",

tflite/testing/BUILD

@@ -6,6 +6,7 @@ load(
     "tf_features_nolayering_check_if_ios",
 )
 load("@org_tensorflow//tensorflow:tensorflow.default.bzl", "pybind_extension")
+load("@org_tensorflow//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
 load(
     "//tflite:build_def.bzl",
     "tflite_custom_android_library",
@@ -170,7 +171,6 @@ py_strict_library(
     name = "zip_test_utils",
     srcs = ["zip_test_utils.py"],
     deps = [
-        ":_pywrap_string_util",
         ":generate_examples_report",
         "//tflite/python:lite",
         "//tflite/tools:flatbuffer_utils",
@@ -179,7 +179,14 @@ py_strict_library(
         "@org_tensorflow//tensorflow/python/framework:convert_to_constants",
         "@org_tensorflow//tensorflow/python/saved_model:signature_constants",
         "@org_tensorflow//third_party/py/numpy",
-    ],
+    ] + if_pywrap(
+        if_false = [
+            ":_pywrap_string_util",
+        ],
+        if_true = [
+            "//tflite/python:pywrap_tflite",
+        ],
+    ),
 )
 
 py_strict_binary(
@@ -698,6 +705,9 @@ pybind_extension(
     additional_stubgen_deps = [
         "@org_tensorflow//third_party/py/numpy:numpy",
     ],
+    common_lib_packages = [
+        "tflite/python",
+    ],
     enable_stub_generation = True,
     features = ["-use_header_modules"],
     pytype_srcs = [

tflite/tools/optimize/python/BUILD

@@ -54,6 +54,9 @@ py_strict_library(
 pybind_extension(
     name = "_pywrap_modify_model_interface",
     srcs = ["modify_model_interface.cc"],
+    common_lib_packages = [
+        "tflite/python",
+    ],
     enable_stub_generation = True,
     pytype_srcs = [
         "_pywrap_modify_model_interface.pyi",

tflite/tools/optimize/sparsity/BUILD

@@ -1,5 +1,6 @@
 load("@org_tensorflow//tensorflow:strict.default.bzl", "py_strict_test")
 load("@org_tensorflow//tensorflow:tensorflow.default.bzl", "pybind_extension")
+load("@org_tensorflow//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["@org_tensorflow//tensorflow:license"],
@@ -12,6 +13,9 @@ package(
 pybind_extension(
     name = "format_converter_wrapper_pybind11",
     srcs = ["format_converter_wrapper_pybind11.cc"],
+    common_lib_packages = [
+        "tflite/python",
+    ],
     copts = [
         "-fexceptions",
         "-fno-strict-aliasing",
@@ -37,5 +41,7 @@ py_strict_test(
         "@absl_py//absl/testing:absltest",
         #internal proto upb dep
         "@org_tensorflow//third_party/py/numpy",
-    ],
+    ] + if_pywrap(
+        if_true = ["//tflite/python:pywrap_tflite"],
+    ),
 )