Internal Build Change

PiperOrigin-RevId: 321387151
Change-Id: I49b2bbc4d78d510c27a226db150599b6fafca6ed

tensorflow/python/BUILD

@@ -3,7 +3,7 @@
 #  ":platform" - Low-level and platform-specific Python code.
 
 load("//tensorflow:tensorflow.bzl", "py_strict_library")
-load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_tpu", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cc_test", "tf_cuda_library", "tf_gen_op_wrapper_py")
+load("//tensorflow:tensorflow.bzl", "cc_header_only_library", "if_mlir", "if_not_windows", "if_tpu", "if_xla_available", "py_test", "py_tests", "tf_cc_shared_object", "tf_cc_test", "tf_cuda_library", "tf_enable_mlir_bridge", "tf_gen_op_wrapper_py")
 
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "tf_monitoring_python_deps")
@@ -5670,6 +5670,12 @@ cc_library(
     ],
 )
 
+py_library(
+    name = "global_test_configuration",
+    deps = if_mlir(["//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_pass_registration"]) +
+           tf_enable_mlir_bridge(),
+)
+
 py_library(
     name = "util",
     srcs = glob(
@@ -5693,6 +5699,10 @@ py_library(
     ],
     deps = [
         ":_pywrap_tf32_execution",
+        # global_test_configuration is added here because all major tests depend on this
+        # library. It isn't possible to add these test dependencies via tensorflow.bzl's
+        # py_test because not all tensorflow tests use tensorflow.bzl's py_test.
+        ":global_test_configuration",
         ":tf_decorator",
         ":tf_export",
         ":tf_stack",
@@ -5702,7 +5712,7 @@ py_library(
         "@six_archive//:six",
         "@wrapt",
         "//tensorflow/tools/compatibility:all_renames_v2",
-    ] + if_mlir(["//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_pass_registration"]),
+    ],
 )
 
 tf_py_test(

tensorflow/tensorflow.bzl

@@ -2145,12 +2145,6 @@ def pywrap_tensorflow_macro(
 # This macro is for running python tests against system installed pip package
 # on Windows.
 #
-# This macro can also enable testing with the experimental mlir bridge when
-# enable_mlir_bridge is true. When it is enabled tests are run both with and without
-# the mlir bridge. Support for enabling the mlir bridge is added here because
-# it allows all tensorflow tests to be configured to be run with and without the
-# mlir bridge.
-#
 # py_test is built as an executable python zip file on Windows, which contains all
 # dependencies of the target. Because of the C++ extensions, it would be very
 # inefficient if the py_test zips all runfiles, plus we don't need them when running
@@ -2168,19 +2162,12 @@ def py_test(deps = [], data = [], kernels = [], **kwargs):
     # Python version placeholder
     if kwargs.get("python_version", None) == "PY3":
         kwargs["tags"] = kwargs.get("tags", []) + ["no_oss_py2"]
-    deps = deps.to_list() if type(deps) == "depset" else deps
     native.py_test(
         # TODO(jlebar): Ideally we'd use tcmalloc here.,
         deps = select({
-                   "//conditions:default": deps,
-                   clean_dep("//tensorflow:no_tensorflow_py_deps"): [],
-               }) +
-               select({
-                   str(Label("//tensorflow:enable_mlir_bridge")): [
-                       "//tensorflow/python:is_mlir_bridge_test_true",
-                   ],
-                   "//conditions:default": [],
-               }),
+            "//conditions:default": deps,
+            clean_dep("//tensorflow:no_tensorflow_py_deps"): [],
+        }),
         data = data + select({
             "//conditions:default": kernels,
             clean_dep("//tensorflow:no_tensorflow_py_deps"): ["//tensorflow/tools/pip_package:win_pip_package_marker"],
@@ -2921,6 +2908,14 @@ def if_mlir(if_true, if_false = []):
         "//conditions:default": if_false,
     })
 
+def tf_enable_mlir_bridge():
+    return select({
+        str(Label("//tensorflow:enable_mlir_bridge")): [
+            "//tensorflow/python:is_mlir_bridge_test_true",
+        ],
+        "//conditions:default": [],
+    })
+
 def if_tpu(if_true, if_false = []):
     """Shorthand for select()ing whether to build for TPUs."""
     return select({

tensorflow/tools/pip_package/pip_smoke_test.py

@@ -75,7 +75,6 @@ PYTHON_TARGETS, PY_TEST_QUERY_EXPRESSION = BuildPyTestDependencies()
 DEPENDENCY_DENYLIST = [
     "//tensorflow/python:extra_py_tests_deps",
     "//tensorflow/cc/saved_model:saved_model_half_plus_two",
-    "//tensorflow:enable_mlir_bridge",
     "//tensorflow:no_tensorflow_py_deps",
     "//tensorflow/tools/pip_package:win_pip_package_marker",
     "//tensorflow/python:test_ops_2",