Use C API to reset PjRt client

PiperOrigin-RevId: 579972556
47-studio-org · Nov 6, 2023 · 9e1ab99 · 9e1ab99
1 parent 9285829
commit 9e1ab99
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/tensorflow/c/experimental/next_pluggable_device/BUILD b/tensorflow/c/experimental/next_pluggable_device/BUILD
@@ -16,7 +16,7 @@ cc_library(
         "//tensorflow/c:c_api_macros_hdrs",
         "//tensorflow/c:kernels_experimental_hdrs",
         "//tensorflow/c:kernels_hdrs",
-        "//tensorflow/c:tf_buffer_internal",
+        "//tensorflow/c:tf_buffer",
         "//tensorflow/c:tf_status_internal",
         "//tensorflow/c:tf_tensor_internal",
         "//tensorflow/compiler/jit:variable_info",

diff --git a/tensorflow/core/common_runtime/next_pluggable_device/BUILD b/tensorflow/core/common_runtime/next_pluggable_device/BUILD
@@ -211,6 +211,7 @@ cc_library(
     name = "c_plugin_op_kernel",
     srcs = ["c_plugin_op_kernel.cc"],
     hdrs = ["c_plugin_op_kernel.h"],
+    copts = ["-DTF_CAPI_WEAK"],
     visibility = ["//visibility:public"],
     deps = [
         ":c_plugin_variable",
@@ -300,7 +301,7 @@ cc_library(
     name = "c_plugin_coordination_service_agent",
     srcs = ["c_plugin_coordination_service_agent.cc"],
     hdrs = ["c_plugin_coordination_service_agent.h"],
-    defines = ["TF_CAPI_WEAK"],
+    copts = ["-DTF_CAPI_WEAK"],
     visibility = ["//visibility:public"],
     deps = [
         ":plugin_coordination_service_agent",
@@ -352,7 +353,7 @@ cc_library(
     name = "c_plugin_variable",
     srcs = ["c_plugin_variable.cc"],
     hdrs = ["c_plugin_variable.h"],
-    defines = ["TF_CAPI_WEAK"],
+    copts = ["-DTF_CAPI_WEAK"],
     visibility = ["//visibility:public"],
     deps = [
         ":plugin_variable",