diff --git a/xla/hlo/ir/BUILD b/xla/hlo/ir/BUILD
index c23eb59be6cf3..ff53f70f3a681 100644
--- a/xla/hlo/ir/BUILD
+++ b/xla/hlo/ir/BUILD
@@ -79,6 +79,7 @@ cc_library(
         "//xla:util",
         "//xla:window_util",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
         "//xla/hlo/parser:hlo_lexer",
         "//xla/service:compilation_environments",
         "//xla/service:computation_layout",
@@ -333,6 +334,7 @@ xla_cc_test(
     deps = [
         ":hlo",
         "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
         "//xla/tsl/platform:test_main",
         "@com_google_benchmark//:benchmark",
         "@com_google_googletest//:gtest",
diff --git a/xla/hlo/ir/hlo_casting_utils_test.cc b/xla/hlo/ir/hlo_casting_utils_test.cc
index 65a48c2c3ab99..e0c0f9afc6688 100644
--- a/xla/hlo/ir/hlo_casting_utils_test.cc
+++ b/xla/hlo/ir/hlo_casting_utils_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace {
diff --git a/xla/hlo/ir/hlo_module.cc b/xla/hlo/ir/hlo_module.cc
index 404d66d8aa3be..c12a90f962005 100644
--- a/xla/hlo/ir/hlo_module.cc
+++ b/xla/hlo/ir/hlo_module.cc
@@ -63,6 +63,7 @@ limitations under the License.
 #include "xla/tsl/platform/status.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
+#include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/fingerprint.h"
 
diff --git a/xla/hlo/ir/hlo_module.h b/xla/hlo/ir/hlo_module.h
index d6ce3d1843eba..17431acd323f3 100644
--- a/xla/hlo/ir/hlo_module.h
+++ b/xla/hlo/ir/hlo_module.h
@@ -54,6 +54,7 @@ limitations under the License.
 #include "xla/status_macros.h"
 #include "xla/tsl/lib/gtl/iterator_range.h"
 #include "xla/tsl/platform/logging.h"
+#include "xla/xla.pb.h"
 
 namespace xla {