Add absl/container:flat_hash_set as dependency for pod_tpu_driver to fix build error
PiperOrigin-RevId: 331044046
Change-Id: I5f53dddd9012fdf0842bf1e57fa39983f6c54091
diff --git a/tensorflow/compiler/xla/python/tpu_driver/BUILD b/tensorflow/compiler/xla/python/tpu_driver/BUILD
index f1b9f6b..70aeb3f 100644
--- a/tensorflow/compiler/xla/python/tpu_driver/BUILD
+++ b/tensorflow/compiler/xla/python/tpu_driver/BUILD
@@ -119,10 +119,11 @@
name = "pod_tpu_driver",
srcs = ["pod_tpu_driver.cc"],
deps = [
- ":tpu_driver",
":grpc_tpu_driver",
+ ":tpu_driver",
":tpu_driver_proto_cc",
"@com_google_absl//absl/container:btree",
+ "@com_google_absl//absl/container:flat_hash_set",
"//tensorflow/compiler/xla/pjrt:semaphore",
"//tensorflow/compiler/xla/pjrt:worker_thread",
"//tensorflow/core:lib",
diff --git a/tensorflow/compiler/xla/python/tpu_driver/pod_tpu_driver.cc b/tensorflow/compiler/xla/python/tpu_driver/pod_tpu_driver.cc
index 352030e..114014e 100644
--- a/tensorflow/compiler/xla/python/tpu_driver/pod_tpu_driver.cc
+++ b/tensorflow/compiler/xla/python/tpu_driver/pod_tpu_driver.cc
@@ -15,6 +15,7 @@
#include "absl/container/btree_map.h"
#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
#include "absl/strings/str_split.h"
#include "tensorflow/compiler/xla/pjrt/semaphore.h"
#include "tensorflow/compiler/xla/pjrt/worker_thread.h"
@@ -37,7 +38,7 @@
class PodEvent : public Event {
public:
- explicit PodEvent(PodTpuDriver* driver, int64 operation_id)
+ explicit PodEvent(PodTpuDriver* driver, int64_t operation_id)
: driver_(driver), operation_id_(operation_id) {}
int64_t operation_id() const { return operation_id_; }
@@ -55,7 +56,7 @@
class CombinedEvent : public PodEvent {
public:
- explicit CombinedEvent(PodTpuDriver* driver, int64 operation_id,
+ explicit CombinedEvent(PodTpuDriver* driver, int64_t operation_id,
std::vector<std::shared_ptr<Event>> events)
: PodEvent(driver, operation_id), events_(events) {}
@@ -629,9 +630,9 @@
std::shared_ptr<::grpc::ChannelCredentials> creds_;
std::vector<std::unique_ptr<TpuDriver>> drivers_;
- std::vector<int32> core_to_driver_id_;
+ std::vector<int32_t> core_to_driver_id_;
std::vector<TpuDriver*> core_to_driver_;
- std::vector<int32> core_to_driver_core_;
+ std::vector<int32_t> core_to_driver_core_;
SystemInfo pod_info_;
absl::Mutex mu_;
