diff --git a/xla/python/BUILD b/xla/python/BUILD
index 284c3068149a0..69fb6ff9e754f 100644
--- a/xla/python/BUILD
+++ b/xla/python/BUILD
@@ -972,6 +972,7 @@ cc_library(
         "//xla/pjrt/c:pjrt_c_api_profiler_extension_hdrs",
         "@tsl//tsl/profiler/lib:profiler_factory",
         "@tsl//tsl/profiler/lib:profiler_interface",
+        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
     ],
 )
 
@@ -1127,6 +1128,7 @@ cc_library(
         "//xla/hlo/transforms:flatten_call_graph",
         "//xla/hlo/transforms:hlo_dce",
         "//xla/hlo/transforms:tuple_simplifier",
+        "//xla/pjrt:compile_options_proto_cc",
         "//xla/pjrt:exceptions",
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt:status_casters",
diff --git a/xla/python/profiler_utils.cc b/xla/python/profiler_utils.cc
index 41b14f62a643a..757390578a3c4 100644
--- a/xla/python/profiler_utils.cc
+++ b/xla/python/profiler_utils.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "xla/pjrt/c/pjrt_c_api_profiler_extension.h"
 #include "tsl/profiler/lib/profiler_factory.h"
 #include "tsl/profiler/lib/profiler_interface.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
 
 namespace xla {
 
diff --git a/xla/python/py_compile_only_client.cc b/xla/python/py_compile_only_client.cc
index 52690075441f9..9dde801ff5a7f 100644
--- a/xla/python/py_compile_only_client.cc
+++ b/xla/python/py_compile_only_client.cc
@@ -74,6 +74,7 @@ limitations under the License.
 #include "xla/tsl/concurrency/ref_count.h"
 #include "xla/tsl/python/lib/core/numpy.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
 #include "tsl/platform/statusor.h"
diff --git a/xla/python/python_ref_manager.cc b/xla/python/python_ref_manager.cc
index b450798e1dff4..6ace28f4a0796 100644
--- a/xla/python/python_ref_manager.cc
+++ b/xla/python/python_ref_manager.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include <deque>
 #include <memory>
 #include <utility>
-#include <vector>
 
 #include "absl/container/inlined_vector.h"
 #include "absl/synchronization/mutex.h"
diff --git a/xla/python/refine_polymorphic_shapes.cc b/xla/python/refine_polymorphic_shapes.cc
index 6ab6a7fcf6436..b4d5b1581f46a 100644
--- a/xla/python/refine_polymorphic_shapes.cc
+++ b/xla/python/refine_polymorphic_shapes.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/python/refine_polymorphic_shapes.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 
diff --git a/xla/python/types.h b/xla/python/types.h
index fece926edd301..59c27d99184e5 100644
--- a/xla/python/types.h
+++ b/xla/python/types.h
@@ -25,6 +25,7 @@ limitations under the License.
 
 #include "absl/container/inlined_vector.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "nanobind/nanobind.h"
 #include "xla/layout.h"
diff --git a/xla/python/xla_compiler.cc b/xla/python/xla_compiler.cc
index 9e397bacc7aa3..8920e010759e0 100644
--- a/xla/python/xla_compiler.cc
+++ b/xla/python/xla_compiler.cc
@@ -63,6 +63,7 @@ limitations under the License.
 #include "xla/layout.h"
 #include "xla/layout_util.h"
 #include "xla/literal.h"
+#include "xla/pjrt/compile_options.pb.h"
 #include "xla/pjrt/exceptions.h"
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/status_casters.h"
diff --git a/xla/python/xplane_to_profile_instructions.cc b/xla/python/xplane_to_profile_instructions.cc
index b0db73556e367..c95bd724f3d5f 100644
--- a/xla/python/xplane_to_profile_instructions.cc
+++ b/xla/python/xplane_to_profile_instructions.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "xla/xla.pb.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/types.h"
+#include "tsl/profiler/protobuf/profiled_instructions.pb.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
 namespace xla {