jax-ml/jax · error · std::runtime_error

Init is not implemented

Error message

Init is not implemented

What it means

The CpuCollectives Python wrapper intentionally stubs out Init because initialization happens at construction/first use of the gloo/mpi collectives object; calling Init() explicitly is an API misuse.

Source

Thrown at jaxlib/jax.cc:377

             absl::StatusOr<std::string> serialized = layout.Serialize();
             xla::ThrowIfError(serialized.status());
             return nb::make_tuple(
                 nb::bytes(serialized->data(), serialized->size()));
           })
      .def("__setstate__", [](xla::PjRtLayout* self, nb::tuple t) {
        nb::bytes serialized = nb::cast<nb::bytes>(t[0]);
        absl::StatusOr<std::shared_ptr<const xla::PjRtLayout>> layout =
            xla::PjRtLayout::Deserialize(
                std::string_view(serialized.c_str(), serialized.size()));
        xla::ThrowIfError(layout.status());
        new (self) xla::PjRtLayout((*layout)->xla_layout());
      });

  nb::class_<xla::cpu::CpuCollectives> cpu_collectives(m, "CpuCollectives");
  cpu_collectives
      .def("Init",
           [](xla::cpu::CpuCollectives*) {
             throw std::runtime_error("Init is not implemented");
           })
      .def("Finalize", [](xla::cpu::CpuCollectives*) {
        throw std::runtime_error("Finalize is not implemented");
      });

  m.def(
      "make_gloo_tcp_collectives",
      [](std::shared_ptr<xla::DistributedRuntimeClient> distributed_client,

         std::optional<std::string> hostname,
         std::optional<std::string> interface)
          -> std::shared_ptr<xla::cpu::CpuCollectives> {
#if defined(__linux__)
        std::shared_ptr<xla::KeyValueStoreInterface> kv_store = nullptr;
        if (distributed_client != nullptr) {
          kv_store = GetDistributedKeyValueStore(distributed_client,
                                                 /*key_prefix=*/"cpu:");
        }

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the Init() call — the constructor initializes the collectives
2. If you need re-initialization, construct a new collectives object via make_gloo_tcp_collectives/make_mpi_collectives

Example fix

# before
coll = jaxlib.make_gloo_tcp_collectives(...)
coll.Init()
# after
coll = jaxlib.make_gloo_tcp_collectives(...)

pass  # do not call Init(); construction initializes

When it happens

Trigger: Calling jaxlib CpuCollectives.Init() from Python, usually by mirroring an older or different collective API that required explicit init.

Common situations: Code ported from other distributed libraries (NCCL-style init/finalize lifecycle) applied to jax's CPU collectives.

Related errors

• Finalize is not implemented
• {name} does not accept integer axis_name. Got axis_name={axe
• {name} wrapped function must be passed at least one argument
• primal and tangent arguments to jax.jvp must be tuples or li
• check_error takes an Error as argument, got type {type(error

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