From 38d156282ecc89509f4b21d80db1a37cb290437a Mon Sep 17 00:00:00 2001
From: Clemens Giuliani <clemens@inailuig.it>
Date: Tue, 19 Mar 2024 19:20:58 +0100
Subject: [PATCH] fix windows build

---
 xla/python/xla.cc | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/xla/python/xla.cc b/xla/python/xla.cc
index 66df4d9336bfa..723445c2e6bb1 100644
--- a/xla/python/xla.cc
+++ b/xla/python/xla.cc
@@ -277,18 +277,19 @@ NB_MODULE(xla_extension, m_nb) {
   nb::class_<cpu::MpiCollectives> mpi_collectives(m_nb, "MpiCollectives", cpu_collectives);
   mpi_collectives.def("Init", &cpu::MpiCollectives::Init);
   mpi_collectives.def("Finalize", &cpu::MpiCollectives::Finalize);
-#endif // _WIN32
-
   m_nb.def(
       "make_mpi_collectives",
       []() -> std::shared_ptr<cpu::MpiCollectives> {
-#ifndef _WIN32
         return std::make_shared<cpu::MpiCollectives>();
+      });
 #else  // _WIN32
+m_nb.def(
+      "make_mpi_collectives",
+      []() -> std::shared_ptr<xla::cpu::CollectivesInterface> {
         throw xla::XlaRuntimeError(
             "make_mpi_collectives is not implemented for Windows");
-#endif // _WIN32
       });
+#endif // _WIN32
 
   m_nb.def(
       "get_tfrt_cpu_client",