diff --git a/xla/BUILD b/xla/BUILD
index c99f1a2a5eb16..1cefa57873e44 100644
--- a/xla/BUILD
+++ b/xla/BUILD
@@ -573,6 +573,7 @@ xla_cc_test(
         ":shape_util",
         ":test",
         ":xla_data_proto_cc",
+        "@com_google_absl//absl/types:span",
         "@tsl//tsl/platform:test_main",
     ],
 )
@@ -639,6 +640,7 @@ xla_cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:logging",
         "@tsl//tsl/platform:macros",
@@ -692,6 +694,7 @@ xla_cc_test(
         ":literal_comparison",
         ":literal_util",
         ":test_helpers",
+        ":xla_data_proto_cc",
         "//xla/tsl/lib/core:status_test_util",
         "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:ml_dtypes",
@@ -720,7 +723,6 @@ cc_library(
         "@eigen_archive//:eigen3",
         "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:logging",
-        "@tsl//tsl/platform:ml_dtypes",
     ],
 )
 
@@ -730,6 +732,7 @@ cc_library(
     hdrs = ["metric_table_report.h"],
     visibility = ["//visibility:public"],
     deps = [
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
diff --git a/xla/index_util.cc b/xla/index_util.cc
index f209be28e7fff..e0e83c3c37fb6 100644
--- a/xla/index_util.cc
+++ b/xla/index_util.cc
@@ -20,7 +20,8 @@ limitations under the License.
 #include <string>
 #include <vector>
 
-#include "absl/strings/str_join.h"
+#include "absl/types/span.h"
+#include "xla/layout_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/util.h"
diff --git a/xla/index_util.h b/xla/index_util.h
index bbafbbac83328..f2ed1be4d8d85 100644
--- a/xla/index_util.h
+++ b/xla/index_util.h
@@ -20,10 +20,14 @@ limitations under the License.
 
 #include <vector>
 
+#include "absl/log/check.h"
+#include "absl/strings/str_join.h"
 #include "absl/types/span.h"
+#include "xla/layout_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/types.h"
+#include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla {
diff --git a/xla/index_util_test.cc b/xla/index_util_test.cc
index de393fadf7b20..333f772f0b4cf 100644
--- a/xla/index_util_test.cc
+++ b/xla/index_util_test.cc
@@ -18,6 +18,8 @@ limitations under the License.
 #include <initializer_list>
 #include <vector>
 
+#include "absl/types/span.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/test.h"
 #include "xla/xla_data.pb.h"
diff --git a/xla/layout.cc b/xla/layout.cc
index d7dcf278797cb..de73698daee48 100644
--- a/xla/layout.cc
+++ b/xla/layout.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
diff --git a/xla/literal_comparison.cc b/xla/literal_comparison.cc
index 42eda4b3fa0aa..fa3a7cda9824c 100644
--- a/xla/literal_comparison.cc
+++ b/xla/literal_comparison.cc
@@ -50,7 +50,6 @@ limitations under the License.
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/logging.h"  // IWYU pragma: keep
-#include "tsl/platform/ml_dtypes.h"
 
 using absl::StrAppend;
 using absl::StrAppendFormat;
diff --git a/xla/literal_comparison_test.cc b/xla/literal_comparison_test.cc
index 893820780276f..04b26922d3add 100644
--- a/xla/literal_comparison_test.cc
+++ b/xla/literal_comparison_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "xla/literal_util.h"
 #include "xla/test_helpers.h"
 #include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/xla_data.pb.h"
 #include "tsl/platform/ml_dtypes.h"
 
 namespace xla {
diff --git a/xla/literal_test.cc b/xla/literal_test.cc
index 42b4340d2ddf8..3ce2f675b1a5a 100644
--- a/xla/literal_test.cc
+++ b/xla/literal_test.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include <gtest/gtest.h>
 #include "absl/base/casts.h"
 #include "absl/hash/hash.h"
 #include "absl/random/random.h"
diff --git a/xla/metric_table_report.cc b/xla/metric_table_report.cc
index e84bf5f9cd637..4b771dd310904 100644
--- a/xla/metric_table_report.cc
+++ b/xla/metric_table_report.cc
@@ -22,10 +22,12 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
 #include "tsl/platform/logging.h"
 
 namespace xla {