jax-ml/jax · error · xla::XlaRuntimeError

make_mpi_collectives is not implemented for Windows

Error message

make_mpi_collectives is not implemented for Windows

What it means

MpiCollectives is only compiled when jaxlib is built with MPI support on non-Windows; the default Windows build registers the function but throws to signal the capability is absent.

Source

Thrown at jaxlib/jax.cc:446

            "make_gloo_tcp_collectives only implemented for linux and macos");
#endif  // defined(__linux__)
      },
      nb::arg("distributed_client"), nb::arg("hostname").none() = std::nullopt,
      nb::arg("interface").none() = std::nullopt);

#if !defined(_WIN32) && !defined(PLATFORM_GOOGLE)
  nb::class_<xla::cpu::MpiCollectives> mpi_collectives(m, "MpiCollectives",
                                                       cpu_collectives);
  mpi_collectives.def("Init", &xla::cpu::MpiCollectives::Init);
  mpi_collectives.def("Finalize", &xla::cpu::MpiCollectives::Finalize);
  m.def("make_mpi_collectives",
        []() -> std::shared_ptr<xla::cpu::MpiCollectives> {
          return std::make_shared<xla::cpu::MpiCollectives>();
        });
#else   // !_WIN32 && !PLATFORM_GOOGLE
  m.def("make_mpi_collectives",
        []() -> std::shared_ptr<xla::cpu::CpuCollectives> {
          throw xla::XlaRuntimeError(
              "make_mpi_collectives is not implemented for Windows");
        });
#endif  // !_WIN32 && !PLATFORM_GOOGLE

  m.def(
      "get_tfrt_cpu_client",
      [](bool asynchronous,
         std::shared_ptr<xla::DistributedRuntimeClient> distributed_client,
         int node_id, int num_nodes,
         std::shared_ptr<xla::cpu::CpuCollectives> collectives,
         std::optional<int> num_devices,
         std::optional<int> get_local_topology_timeout_minutes,
         std::optional<int> get_global_topology_timeout_minutes,
         std::optional<xla::ifrt::TransferServerInterfaceFactory>
             transfer_server_factory) -> nb_class_ptr<PyClient> {
        std::unique_ptr<xla::ifrt::PjRtClient> ifrt_client;
        {
          nb::gil_scoped_release gil_release;

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use gloo collectives (make_gloo_tcp_collectives) instead where supported
Build jaxlib from source with MPI on Linux
Run under WSL2/Linux with an MPI-enabled build

Example fix

# before
coll = jaxlib.make_mpi_collectives()
# after
coll = jaxlib.make_gloo_tcp_collectives(distributed_client)

Defensive patterns

Strategy: fallback

Validate before calling

import sys
mpi_ok = sys.platform != 'win32' and hasattr(jaxlib, 'make_mpi_collectives')

Type guard

def mpi_available() -> bool:
    import jaxlib
    return sys.platform != 'win32' and getattr(jaxlib, '_mpi_support', False)

Try / catch

try:
    coll = jaxlib.make_mpi_collectives()
except Exception:
    coll = jaxlib.make_gloo_tcp_collectives(client)

Prevention

Prefer gloo collectives unless you built jaxlib with MPI

When it happens

Trigger: Calling jaxlib.make_mpi_collectives() on Windows or on a jaxlib built without MPI support.

Common situations: Attempting MPI-based CPU collectives on Windows or with the standard pip jaxlib (which lacks MPI).

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/6380cfcef6e302ed. Report an issue: GitHub.