{"record":{"id":"d745bc313bc7b776","repo":"jax-ml/jax","slug":"cannot-register-xla-transform-s-pjrt-plugin-do","errorCode":null,"errorMessage":"Cannot register XLA transform '%s': PJRT plugin does not support the XlaTransform extension.","messagePattern":"Cannot register XLA transform '(.+?)': PJRT plugin does not support the XlaTransform extension\\.","errorType":"exception","errorClass":"runtime_error","httpStatus":null,"severity":"error","filePath":"jaxlib/xla.cc","lineNumber":233,"sourceCode":"  // 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\n        // clear_xla_transform is called.\n        auto* state = new CApiCallbackState();\n        state->py_callback = std::move(callback);\n        state->callbacks.version = PJRT_API_XLA_TRANSFORM_EXTENSION_VERSION;\n        state->callbacks.dtor = [](PJRT_XlaTransform_Callbacks* callbacks) {\n          auto* state = reinterpret_cast<CApiCallbackState*>(\n              reinterpret_cast<char*>(callbacks) -\n              offsetof(CApiCallbackState, callbacks));\n          delete state;\n        };\n        state->callbacks.transform_hlo_module = CApiTransformHloModuleCallback;\n","sourceCodeStart":215,"sourceCodeEnd":251,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/xla.cc#L215-L251","documentation":"When registering an XLA transform through a PJRT plugin's C API, jaxlib looks up the PJRT_Xla_Transform_Extension. If the loaded plugin doesn't implement this extension, registration fails with this runtime error naming the transform.","triggerScenarios":"Calling register_xla_transform_c_api with a valid pjrt_c_api capsule from a plugin (e.g. a CPU backend or older TPU/GPU plugin) that doesn't provide the XlaTransform extension.","commonSituations":"Testing plugin integration code against the CPU backend or jax's built-in backends, which lack the extension; using an outdated plugin version that predates XlaTransform support.","solutions":["Upgrade the plugin (TPU/GPU plugin) to a version advertising the XlaTransform extension","Skip registration when the extension is absent (feature-detect before registering)","Run transform registration only against backends known to support it"],"exampleFix":"# before\nregister_xla_transform_c_api(capsule, 'my_pass', 0, cb)  # CPU plugin\n\n# after\nif getattr(backend.client._plugin, 'supports_xla_transform', False):\n    register_xla_transform_c_api(capsule, 'my_pass', 0, cb)","handlingStrategy":"fallback","validationCode":"# feature-detect before registering\napi = plugin._C_API\n# cheap probe: registration is the only check; guard with try/except in fallback","typeGuard":null,"tryCatchPattern":"try:\n    register_xla_transform_c_api(api, name, stage, cb)\nexcept RuntimeError as e:\n    if 'XlaTransform extension' in str(e):\n        logging.warning('plugin %s lacks XlaTransform; skipping transform', name)\n    else:\n        raise","preventionTips":["Only run transform registration against backends that advertise the extension","Upgrade TPU/GPU plugins before using XLA transform registration","Skip gracefully in CPU test environments"],"tags":["jax","xla","pjrt","plugin-capabilities"],"backgroundTag":"unsupported-plugin-extension","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}