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

Finalize is not implemented

Error message

Finalize is not implemented

What it means

CpuCollectives.Finalize() is a deliberate stub; teardown is handled by the C++ destructor / garbage collection of the collectives object, not by an explicit Python call.

Source

Thrown at jaxlib/jax.cc:380

                 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:");
        }
        auto gloo_kv_store =
            std::make_unique<xla::cpu::GlooKeyValueStore>(kv_store);
        auto tcp_attrs = gloo::transport::tcp::attr();

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Delete the Finalize() call; rely on object deallocation
2. For deterministic teardown, del the collectives reference and force gc.collect()

Example fix

# before
try:
    coll.Finalize()
finally:
    pass
# after
# no explicit finalize needed

pass  # do not call Finalize(); GC handles teardown

When it happens

Trigger: Calling CpuCollectives.Finalize() explicitly at shutdown in Python.

Common situations: Distributed scripts adding symmetric init/finalize cleanup calls borrowed from MPI-style APIs.

Related errors

• Init 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/c11d3fbdc9a99a3d. Report an issue: GitHub.