From db4fd942efa479cf2bb9f77916d50b03e4ef5319 Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Sat, 1 Feb 2025 17:32:18 -0800
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 722211370
---
 xla/pjrt/BUILD                 | 5 +++++
 xla/pjrt/host_callback.cc      | 1 +
 xla/pjrt/host_callback.h       | 1 +
 xla/pjrt/host_callback_test.cc | 5 +++++
 4 files changed, 12 insertions(+)

diff --git a/xla/pjrt/BUILD b/xla/pjrt/BUILD
index 5502da86e87ab..8e4ed5cbe0d3c 100644
--- a/xla/pjrt/BUILD
+++ b/xla/pjrt/BUILD
@@ -954,6 +954,7 @@ cc_library(
         ":pjrt_future",
         "//xla:shape_util",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
@@ -968,9 +969,13 @@ xla_cc_test(
     deps = [
         ":host_callback",
         ":pjrt_client",
+        "//xla:xla_data_proto_cc",
         "//xla/tests:literal_test_util",
         "//xla/tsl/lib/core:status_test_util",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/synchronization",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/xla/pjrt/host_callback.cc b/xla/pjrt/host_callback.cc
index 1faa844936cc5..51f079b067eb7 100644
--- a/xla/pjrt/host_callback.cc
+++ b/xla/pjrt/host_callback.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "xla/pjrt/pjrt_client.h"
diff --git a/xla/pjrt/host_callback.h b/xla/pjrt/host_callback.h
index 2777cd368d010..9b32f7daf677e 100644
--- a/xla/pjrt/host_callback.h
+++ b/xla/pjrt/host_callback.h
@@ -25,6 +25,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/base/thread_annotations.h"
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
diff --git a/xla/pjrt/host_callback_test.cc b/xla/pjrt/host_callback_test.cc
index ef9d5d9ec70c5..b6722892afb92 100644
--- a/xla/pjrt/host_callback_test.cc
+++ b/xla/pjrt/host_callback_test.cc
@@ -15,15 +15,20 @@ limitations under the License.
 
 #include "xla/pjrt/host_callback.h"
 
+#include <cstdint>
 #include <cstring>
 #include <memory>
 #include <utility>
 
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/synchronization/notification.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/tests/literal_test_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace {