jax-ml/jax · error · nb::value_error

c_api argument to load_pjrt_plugin is not a pjrt_c_api capsu

Error message

c_api argument to load_pjrt_plugin is not a pjrt_c_api capsule.

What it means

When load_pjrt_plugin is called with an existing c_api capsule (rather than a library path), that capsule must be named 'pjrt_c_api' so jaxlib can trust its payload is a PJRT_Api*.

Source

Thrown at jaxlib/jax.cc:522

      nb::arg("num_devices").none() = std::nullopt,
      nb::arg("get_local_topology_timeout_minutes").none() = std::nullopt,
      nb::arg("get_global_topology_timeout_minutes").none() = std::nullopt,
      nb::arg("transfer_server_factory").none() = std::nullopt);
  m.def("pjrt_plugin_loaded", [](std::string platform_name) -> bool {
    absl::StatusOr<const PJRT_Api*> pjrt_api = pjrt::PjrtApi(platform_name);
    return pjrt_api.ok();
  });
  m.def(
      "load_pjrt_plugin",
      [](std::string platform_name, std::optional<std::string> library_path,
         std::optional<nb::capsule> c_api) -> nb::capsule {
        if (library_path.has_value()) {
          const PJRT_Api* api = xla::ValueOrThrow(
              pjrt::LoadPjrtPlugin(platform_name, *library_path));
          return nb::capsule(absl::bit_cast<void*>(api), "pjrt_c_api");
        }
        if (std::string_view(c_api->name()) != "pjrt_c_api") {
          throw nb::value_error(
              "c_api argument to load_pjrt_plugin is not a pjrt_c_api "
              "capsule.");
        }
        xla::ThrowIfError(pjrt::SetPjrtApi(
            platform_name, static_cast<const PJRT_Api*>(c_api->data())));
        return *c_api;
      },
      nb::arg("platform_name"), nb::arg("library_path").none() = std::nullopt,
      nb::arg("c_api").none() = std::nullopt);
  m.def(
      "get_pjrt_plugin",
      [](std::string platform_name) -> nb::capsule {
        const PJRT_Api* api = xla::ValueOrThrow(pjrt::PjrtApi(platform_name));
        return nb::capsule(absl::bit_cast<void*>(api), "pjrt_c_api");
      },
      nb::arg("platform_name"));
  m.def("pjrt_plugin_initialized", [](std::string platform_name) -> bool {
    return xla::ValueOrThrow(pjrt::IsPjrtPluginInitialized(platform_name));

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Call load_pjrt_plugin(platform_name, library_path=...) to let jaxlib create the capsule
  2. Pass the exact capsule previously returned by load_pjrt_plugin or get_backend_c_api

Example fix

# before
jaxlib.load_pjrt_plugin('mydev', c_api=mystery_capsule)
# after
api = jaxlib.load_pjrt_plugin('mydev', library_path='/path/libplugin.so')
Defensive patterns

Strategy: validation

Validate before calling

assert c_api is None or getattr(c_api, 'name', lambda: None)() == 'pjrt_c_api'

Prevention

When it happens

Trigger: jaxlib.load_pjrt_plugin(platform, c_api=cap) where cap came from anything other than a prior load_pjrt_plugin result or a properly named capsule.

Common situations: Plugin bootstrap code that creates its own capsule with a different name, or passes a capsule that was consumed/renamed.

Related errors


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