From baf117060313ab3141f4f6fbdb00e70a310473a5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 12 Sep 2024 22:28:29 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 674147797
---
 xla/python/ifrt_proxy/server/BUILD           | 3 ++-
 xla/python/ifrt_proxy/server/grpc_server.cc  | 1 +
 xla/python/ifrt_proxy/server/host_buffer.h   | 1 +
 xla/python/ifrt_proxy/server/ifrt_backend.cc | 1 -
 4 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/xla/python/ifrt_proxy/server/BUILD b/xla/python/ifrt_proxy/server/BUILD
index ae8f52238db2d..9eb86aa8308b1 100644
--- a/xla/python/ifrt_proxy/server/BUILD
+++ b/xla/python/ifrt_proxy/server/BUILD
@@ -29,6 +29,7 @@ cc_library(
         "//xla/python/ifrt",
         "//xla/python/ifrt_proxy/common:grpc_credentials",
         "//xla/python/ifrt_proxy/common:grpc_ifrt_service_cc_grpc_proto",
+        "//xla/python/ifrt_proxy/common:ifrt_service_proto_cc",
         "@com_github_grpc_grpc//:grpc",
         "@com_github_grpc_grpc//:grpc++",
         "@com_google_absl//absl/functional:any_invocable",
@@ -122,7 +123,6 @@ cc_library(
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
-        "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_layout",
         "//xla/python/ifrt",
         "//xla/python/ifrt:program_serdes",
@@ -226,6 +226,7 @@ cc_library(
     srcs = ["host_buffer.cc"],
     hdrs = ["host_buffer.h"],
     deps = [
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/xla/python/ifrt_proxy/server/grpc_server.cc b/xla/python/ifrt_proxy/server/grpc_server.cc
index 6295dc11c97d8..936cd96443d97 100644
--- a/xla/python/ifrt_proxy/server/grpc_server.cc
+++ b/xla/python/ifrt_proxy/server/grpc_server.cc
@@ -32,6 +32,7 @@
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt_proxy/common/grpc_credentials.h"
 #include "xla/python/ifrt_proxy/common/grpc_ifrt_service.grpc.pb.h"
+#include "xla/python/ifrt_proxy/common/ifrt_service.pb.h"
 #include "xla/python/ifrt_proxy/server/grpc_service_impl.h"
 #include "xla/python/ifrt_proxy/server/host_buffer.h"
 #include "xla/python/ifrt_proxy/server/ifrt_backend.h"
diff --git a/xla/python/ifrt_proxy/server/host_buffer.h b/xla/python/ifrt_proxy/server/host_buffer.h
index f9b07a40f30e9..0d82898be9258 100644
--- a/xla/python/ifrt_proxy/server/host_buffer.h
+++ b/xla/python/ifrt_proxy/server/host_buffer.h
@@ -20,6 +20,7 @@
 #include <memory>
 #include <string>
 
+#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
diff --git a/xla/python/ifrt_proxy/server/ifrt_backend.cc b/xla/python/ifrt_proxy/server/ifrt_backend.cc
index 9b643db070f2a..040c0444cbd20 100644
--- a/xla/python/ifrt_proxy/server/ifrt_backend.cc
+++ b/xla/python/ifrt_proxy/server/ifrt_backend.cc
@@ -40,7 +40,6 @@
 #include "absl/types/span.h"
 #include "llvm/Support/Casting.h"
 #include "xla/layout.h"
-#include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/compiler.h"