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

c_api argument to clear_xla_transform_c_api is not a pjrt_c_

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 mirrors the registration check: its first argument must be a PyCapsule named 'pjrt_c_api'. Passing any other capsule or object raises this value error.

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)

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.

Related errors

• c_api argument to register_xla_transform_c_api is not a pjrt
• Cannot register XLA transform '%s': PJRT plugin does not sup
• multi-platform lowering for buffer_callback
• Arguments to after_all must be tokens
• Argument to register_custom_call_partitioner was not a pjrt_

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