diff --git a/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc b/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc
index ed71770ce69e..2912ffd72566 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc
+++ b/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc
@@ -41,8 +41,6 @@ limitations under the License.
 #include "jaxlib/mosaic/dialect/tpu/tpu_enums.cc.inc"
 #include "xla/layout.h"
 
-// This is a bit unclean, but we need to squat the xla namespace to make sure
-// that this overload is found via argument-dependent lookup.
 namespace xla {
 
 llvm::hash_code hash_value(const ::xla::Tile &p) { return absl::HashOf(p); }
diff --git a/jaxlib/mosaic/dialect/tpu/tpu_dialect.h b/jaxlib/mosaic/dialect/tpu/tpu_dialect.h
index b6c4434ccc2d..382aaa72f15d 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu_dialect.h
+++ b/jaxlib/mosaic/dialect/tpu/tpu_dialect.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <memory>
 #include <utility>
 
+#include "llvm/ADT/Hashing.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -31,6 +32,12 @@ limitations under the License.
 #include "jaxlib/mosaic/dialect/tpu/tpu_enums.h.inc"
 #include "xla/layout.h"  // IWYU pragma: keep
 
+// This is a bit unclean, but we need to squat the xla namespace to make sure
+// that this overload is found via argument-dependent lookup.
+namespace xla {
+llvm::hash_code hash_value(const xla::Tile& tile);
+}
+
 namespace mlir::tpu {
 class TPUDialect;
 }  // namespace mlir::tpu