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

c_api argument to clear_xla_transform_c_api is not a…

Error message

c_api argument to clear_xla_transform_c_api is not a pjrt_c_api capsule.

What it means

clear_xla_transform_c_api retrieves the plugin's API table from a pjrt_c_api capsule before deregistering transforms. It fires when the c_api argument is not such a capsule, so the extension table cannot be located.

Solutions

1. Pass the same pjrt_c_api capsule used at registration (plugin._C_API)
2. Store the capsule reference used for registration and reuse it in clear

Example fix

# before
clear_xla_transform_c_api(backend_client, 'my_pass', 0)

# after
clear_xla_transform_c_api(plugin._C_API, 'my_pass', 0)  # 'pjrt_c_api' capsule

api = plugin._C_API  # the 'pjrt_c_api' capsule used at registration
if type(api).__name__ != 'PyCapsule':
    raise TypeError('expected pjrt_c_api capsule')

def is_pjrt_capsule(c) -> bool:
    return type(c).__name__ == 'PyCapsule'

When it happens

Trigger: Calling jaxlib's clear_xla_transform_c_api with the wrong capsule (different name) or a non-capsule object.

Common situations: Cleanup code that passes the backend client or module instead of the plugin's _C_API capsule; capsule name lost during serialization/handoff.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• c_api argument to register_xla_transform_c_api is not a…
• Cannot register XLA transform
• Argument to get_c_api_topology contained a null pointer.
• Argument to get_c_api_topology was not a pjrt_c_api capsule.
• Argument to register_custom_call_partitioner was not a…

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

Appendix: source

Thrown at jaxlib/xla.cc:272

        args.name = name.c_str();
        args.name_size = name.size();
        args.stage = static_cast<PJRT_XlaTransform_PipelineStage>(stage);
        args.callbacks = &state->callbacks;

        PJRT_Error* error = extension->register_xla_transform(&args);
        if (error != nullptr) {
          absl::Status status = pjrt::PjrtErrorToStatus(error);
          throw std::runtime_error(status.ToString());
        }
      },
      nb::arg("c_api"), nb::arg("name"), nb::arg("stage"),
      nb::arg("callback"));

  m.def(
      "clear_xla_transform_c_api",
      [](nb::capsule c_api, std::string name, int stage) {
        if (std::string_view(c_api.name()) != "pjrt_c_api") {
          throw nb::value_error(
              "c_api argument to clear_xla_transform_c_api is not a "
              "pjrt_c_api capsule.");
        }
        const PJRT_Api* c_api_value =
            static_cast<const PJRT_Api*>(c_api.data());

        PJRT_Xla_Transform_Extension* extension =
            pjrt::FindExtension<PJRT_Xla_Transform_Extension>(
                c_api_value,
                PJRT_Extension_Type::PJRT_Extension_Type_XlaTransform);
        if (extension == nullptr) {
          return false;
        }

        PJRT_Clear_Xla_Transform_Args args;
        args.struct_size = PJRT_Clear_Xla_Transform_Args_STRUCT_SIZE;
        args.name = name.c_str();
        args.name_size = name.size();

View on GitHub (pinned to 1e1c6a8fc0)