diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc
index 62337c670af573..fa3dbfe2ee65ca 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc
@@ -81,6 +81,7 @@ limitations under the License.
 #include "xla/status_macros.h"
 #include "xla/tsl/concurrency/ref_count.h"
 #include "xla/util.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/platform/casts.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.cc
index 988789bf69eccd..9ecdf64b4a72a0 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.cc
@@ -17,6 +17,8 @@ limitations under the License.
 
 #include <string>
 
+#include "absl/status/statusor.h"
+
 namespace xla {
 namespace ifrt {
 
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.h
index b192cbf76144ef..8c539af83e9930 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_host_callback.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/status/statusor.h"
 #include "xla/pjrt/host_callback.h"
 #include "xla/python/ifrt/host_callback.h"
 
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_memory.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_memory.h
index 1d01683fe05369..3964ac56b184d5 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_memory.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_memory.h
@@ -16,7 +16,10 @@ limitations under the License.
 #ifndef XLA_PYTHON_PJRT_IFRT_PJRT_MEMORY_H_
 #define XLA_PYTHON_PJRT_IFRT_PJRT_MEMORY_H_
 
+#include <vector>
+
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "llvm/Support/ExtensibleRTTI.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/python/ifrt/memory.h"
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.cc
index d27097145828a2..9c5eb61100cd52 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_device_description.h"
 #include "xla/python/ifrt/attribute_map.h"
 #include "xla/python/pjrt_ifrt/pjrt_attribute_map_util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla::ifrt {
 
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.h
index 82fc59c8005c01..81adf1bda215df 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_topology.h
@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_device_description.h"
 #include "xla/python/ifrt/attribute_map.h"
 #include "xla/python/ifrt/topology.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla::ifrt {
 
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc
index 17e529e7695650..b863828486892d 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.cc
@@ -15,12 +15,17 @@ limitations under the License.
 
 #include "xla/python/pjrt_ifrt/pjrt_tuple.h"
 
-#include <memory>
-#include <optional>
 #include <string>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
 #include "llvm/Support/ExtensibleRTTI.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h
index 2a6359be9507b2..2140d19473cc51 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_tuple.h
@@ -22,6 +22,9 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/synchronization/notification.h"
 #include "absl/types/span.h"
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h
index 0fefd73cf2b0cf..1cf81e253da86e 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/xla_compiler.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/status/statusor.h"
 #include "llvm/Support/ExtensibleRTTI.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_serdes_test.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_serdes_test.cc
index 6a3126e941f2cb..2474b51e3a50b7 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_serdes_test.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_serdes_test.cc
@@ -13,7 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <cstdint>
 #include <memory>
 
 #include <gmock/gmock.h>