diff --git a/xla/backends/cpu/BUILD b/xla/backends/cpu/BUILD
index c65516d819b73..726b57671f6e8 100644
--- a/xla/backends/cpu/BUILD
+++ b/xla/backends/cpu/BUILD
@@ -114,6 +114,7 @@ cc_library(
         "//xla:literal",
         "//xla:shape_util",
         "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service:buffer_assignment",
         "//xla/stream_executor:device_memory",
diff --git a/xla/backends/cpu/constant_allocation.cc b/xla/backends/cpu/constant_allocation.cc
index 9da4f81481098..3702811c48e09 100644
--- a/xla/backends/cpu/constant_allocation.cc
+++ b/xla/backends/cpu/constant_allocation.cc
@@ -35,6 +35,7 @@ limitations under the License.
 #include "xla/stream_executor/device_memory.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla::cpu {