Enable unified_api_test in OSS.

Needed to make sure that the pybind targets have header-only dependencies.

PiperOrigin-RevId: 330925715
Change-Id: I9b63049cd5a02af73a764837c954586ab7da8af0

tensorflow/c/eager/BUILD

@@ -101,6 +101,7 @@ filegroup(
         "c_api_experimental.h",
         "c_api_internal.h",
         "c_api_unified_experimental.h",
+        "c_api_unified_experimental_internal.h",
         "dlpack.h",
         "immediate_execution_context.h",
         "immediate_execution_operation.h",

tensorflow/c/experimental/ops/BUILD

@@ -65,3 +65,28 @@ cc_library(
         "//tensorflow/core/platform:errors",
     ],
 )
+
+cc_library(
+    name = "ops",
+    visibility = [
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":array_ops",
+        ":math_ops",
+        ":nn_ops",
+    ],
+)
+
+filegroup(
+    name = "pywrap_required_hdrs",
+    srcs = [
+        "array_ops.h",
+        "math_ops.h",
+        "nn_ops.h",
+    ],
+    visibility = [
+        "//tensorflow/core:__pkg__",
+        "//tensorflow/python:__pkg__",
+    ],
+)

tensorflow/python/BUILD

@@ -5938,6 +5938,8 @@ pywrap_tensorflow_macro(
         "//tensorflow/c:tf_status_helper",
         "//tensorflow/c/eager:c_api",
         "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/c/experimental/ops",
+        "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
         "//tensorflow/core/data/service:server_lib",
         "//tensorflow/core/distributed_runtime/rpc:grpc_rpc_factory_registration",
         "//tensorflow/core/distributed_runtime/rpc:grpc_server_lib",
@@ -7431,6 +7433,36 @@ tf_python_pybind_extension(
     ],
 )
 
+cc_library(
+    name = "unified_api_pywrap_required_headers",
+    textual_hdrs = [
+        "lib/core/numpy.h",
+        "lib/core/py_exception_registry.h",
+        "lib/core/pybind11_status.h",
+        "lib/core/bfloat16.h",
+        "lib/core/ndarray_tensor.h",
+        "lib/core/ndarray_tensor_bridge.h",
+        "lib/core/safe_ptr.h",
+        "lib/core/safe_pyobject_ptr.h",
+        "//tensorflow/c:headers",
+        "//tensorflow/c/eager:headers",
+        "//tensorflow/c/eager:pywrap_required_hdrs",
+        "//tensorflow/c/experimental/ops:pywrap_required_hdrs",
+        "//tensorflow/core/common_runtime/eager:pywrap_required_hdrs",
+        "//tensorflow/python/eager:pywrap_required_hdrs",
+        "//tensorflow/core/common_runtime:core_cpu_lib_headers",
+        "//tensorflow/core/public:session.h",
+        "//tensorflow/core/public:session_options.h",
+    ],
+    visibility = ["//tensorflow/python/framework/experimental:__pkg__"],
+    deps = [
+        "//tensorflow/c:pywrap_required_hdrs",
+        "//tensorflow/core/framework:pywrap_required_hdrs",
+        "//third_party/py/numpy:headers",
+        "//third_party/python_runtime:headers",
+    ],
+)
+
 py_library(
     name = "pywrap_tfe",
     srcs = ["pywrap_tfe.py"],
@@ -7445,6 +7477,7 @@ tf_python_pybind_extension(
     name = "_pywrap_tfe",
     srcs = ["tfe_wrapper.cc"],
     hdrs = [
+        "lib/core/numpy.h",
         "lib/core/safe_ptr.h",
         "util/util.h",
         ":py_exception_registry_hdr",
@@ -7460,6 +7493,7 @@ tf_python_pybind_extension(
     deps = [
         ":safe_pyobject_ptr",
         ":pybind11_lib",
+        "//third_party/py/numpy:headers",
         ":pybind11_status",
         "//tensorflow/core/framework:pywrap_required_hdrs",
         "@com_google_absl//absl/container:flat_hash_map",

tensorflow/python/eager/BUILD

@@ -113,6 +113,7 @@ cuda_py_test(
 filegroup(
     name = "pywrap_required_hdrs",
     srcs = [
+        "pywrap_tensor.h",
         "pywrap_tensor_conversion.h",
         "pywrap_tfe.h",
     ],

tensorflow/python/framework/experimental/BUILD

@@ -17,24 +17,12 @@ tf_python_pybind_extension(
     features = ["-layering_check"],
     module_name = "_unified_api",
     deps = [
-        "//tensorflow/c/eager:abstract_context",
-        "//tensorflow/c/eager:abstract_function",
-        "//tensorflow/c/eager:abstract_operation",
-        "//tensorflow/c/eager:abstract_tensor_handle",
-        "//tensorflow/c/eager:c_api",
-        "//tensorflow/c/eager:c_api_experimental",
-        "//tensorflow/c/eager:c_api_unified_internal",
-        "//tensorflow/c/eager:immediate_execution_context",
-        "//tensorflow/c/eager:immediate_execution_tensor_handle",
         "//tensorflow/c/eager:tfe_tensorhandle_internal",
-        "//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/lib/llvm_rtti",
         "//tensorflow/python:pybind11_lib",
-        "//tensorflow/python:pybind11_status",
-        "//tensorflow/python:safe_ptr",
-        "//tensorflow/python/eager:pywrap_tfe_lib",
+        "//tensorflow/python:unified_api_pywrap_required_headers",
         "@pybind11",
     ],
 )
@@ -44,11 +32,13 @@ tf_python_pybind_extension(
     srcs = ["math_ops.cc"],
     module_name = "_math_ops",
     deps = [
-        "//tensorflow/c/eager:abstract_context",
-        "//tensorflow/c/eager:abstract_tensor_handle",
-        "//tensorflow/c/experimental/ops:math_ops",
+        "//tensorflow/c/eager:tfe_tensorhandle_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/lib/llvm_rtti",
         "//tensorflow/python:pybind11_lib",
-        "//tensorflow/python:pybind11_status",
+        "//tensorflow/python:unified_api_pywrap_required_headers",
         "@com_google_absl//absl/types:span",
         "@pybind11",
     ],
@@ -75,8 +65,10 @@ cuda_py_test(
     size = "small",
     srcs = ["unified_api_test.py"],
     tags = [
-        "no_oss",  # TODO(srbs): Enable.
-        "no_pip",  # TODO(srbs): Enable.
+        # Note(srbs): These python bindings are not
+        # exported as part of the pip package yet so
+        # this test is disabled.
+        "no_pip",
     ],
     deps = [
         ":_unified_api",