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
- Call load_pjrt_plugin(platform_name, library_path=...) to let jaxlib create the capsule
- 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
- Use library_path mode when in doubt
- Only forward capsules produced by load_pjrt_plugin
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
- Argument to get_c_api_topology was not a pjrt_c_api capsule.
- Argument to register_custom_call_partitioner was not a pjrt_
- Argument to get_c_api_topology contained a null pointer.
- {full_name} must be a pytree prefix with bool leaves or a tu
- Axes mentioned in `manual_axis_type` field of ShapedArray sh
AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27).
Data as JSON: /api/errors/3c528027340815f7.
Report an issue: GitHub.