{"record":{"id":"0c034f9452adee84","repo":"jax-ml/jax","slug":"c-api-argument-to-register-xla-transform-c-api-is","errorCode":null,"errorMessage":"c_api argument to register_xla_transform_c_api is not a pjrt_c_api capsule.","messagePattern":"c_api argument to register_xla_transform_c_api is not a pjrt_c_api capsule\\.","errorType":"validation","errorClass":"nb::value_error","httpStatus":null,"severity":"error","filePath":"jaxlib/xla.cc","lineNumber":221,"sourceCode":"            break;\n          case 1:\n            pipeline_stage =\n                xla::HloXlaTransform::PipelineStage::kPostScheduler;\n            break;\n          default:\n            throw std::runtime_error(\"Invalid pipeline stage\");\n        }\n        return xla::ClearHloXlaTransform(pipeline_stage, name);\n      },\n      nb::arg(\"name\"), nb::arg(\"stage\"));\n\n  // Register a transform via the PJRT C API XlaTransform extension.\n  // This is used for plugin backends (e.g. TPU, GPU).\n  m.def(\n      \"register_xla_transform_c_api\",\n      [](nb::capsule c_api, std::string name, int stage, nb::object callback) {\n        if (std::string_view(c_api.name()) != \"pjrt_c_api\") {\n          throw nb::value_error(\n              \"c_api argument to register_xla_transform_c_api is not a \"\n              \"pjrt_c_api capsule.\");\n        }\n        const PJRT_Api* c_api_value =\n            static_cast<const PJRT_Api*>(c_api.data());\n\n        PJRT_Xla_Transform_Extension* extension =\n            pjrt::FindExtension<PJRT_Xla_Transform_Extension>(\n                c_api_value,\n                PJRT_Extension_Type::PJRT_Extension_Type_XlaTransform);\n        if (extension == nullptr) {\n          throw std::runtime_error(\n              absl::StrCat(\"Cannot register XLA transform '\", name,\n                           \"': PJRT plugin does not support the XlaTransform \"\n                           \"extension.\"));\n        }\n\n        // Allocate callback state on the heap. Cleared via dtor if","sourceCodeStart":203,"sourceCodeEnd":239,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/xla.cc#L203-L239","documentation":"register_xla_transform_c_api expects the first argument to be a PyCapsule whose name is exactly 'pjrt_c_api' (the capsule plugin backends expose). Any other capsule name or type raises this value error.","triggerScenarios":"Calling jaxlib's register_xla_transform_c_api with a capsule obtained from a different source (e.g. a 'xla_extension' capsule), a plain object, or a capsule whose name was stripped.","commonSituations":"Plugin/backend integration code that grabs the wrong capsule from the PJRT plugin module; passing the module object instead of its _C_API attribute.","solutions":["Pass the plugin's pjrt_c_api capsule: typically plugin_module._C_API (named 'pjrt_c_api')","Verify with capsule name inspection before calling","Use jax.lib.xla_bridge / plugin APIs to obtain the correct capsule"],"exampleFix":"# before\nregister_xla_transform_c_api(jax.lib.xla_bridge.get_backend().client, ...)\n\n# after\nimport jax._src.lib.xla_client as xc\ncapsule = jax.lib.xla_bridge.get_backend().client._plugin._C_API  # name 'pjrt_c_api'\nregister_xla_transform_c_api(capsule, 'pass', 0, cb)","handlingStrategy":"type-guard","validationCode":"import ctypes\n# nanobind capsule; check name via its .name attribute if exposed, or plugin attribute\napi = getattr(plugin_module, '_C_API', None)\nif api is None:\n    raise ValueError('plugin does not expose _C_API pjrt capsule')","typeGuard":"def is_pjrt_capsule(c) -> bool:\n    return type(c).__name__ == 'PyCapsule' and getattr(c, 'name', lambda: None)() == 'pjrt_c_api' if False else type(c).__name__ == 'PyCapsule'","tryCatchPattern":"try:\n    register_xla_transform_c_api(capsule, name, stage, cb)\nexcept ValueError as e:\n    if 'pjrt_c_api capsule' in str(e):\n        capsule = plugin._C_API\n        register_xla_transform_c_api(capsule, name, stage, cb)\n    else:\n        raise","preventionTips":["Always source the capsule from plugin._C_API","Keep a single reference to the capsule for both register and clear"],"tags":["jax","xla","pjrt","pycapsule"],"backgroundTag":"wrong-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}