jax-ml/jax · error · runtime_error

Cannot register XLA transform '%s': PJRT plugin does not sup

Error message

Cannot register XLA transform '%s': PJRT plugin does not support the XlaTransform extension.

What it means

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.

Source

Thrown at jaxlib/xla.cc:233

  // Register a transform via the PJRT C API XlaTransform extension.
  // This is used for plugin backends (e.g. TPU, GPU).
  m.def(
      "register_xla_transform_c_api",
      [](nb::capsule c_api, std::string name, int stage, nb::object callback) {
        if (std::string_view(c_api.name()) != "pjrt_c_api") {
          throw nb::value_error(
              "c_api argument to register_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) {
          throw std::runtime_error(
              absl::StrCat("Cannot register XLA transform '", name,
                           "': PJRT plugin does not support the XlaTransform "
                           "extension."));
        }

        // Allocate callback state on the heap. Cleared via dtor if
        // clear_xla_transform is called.
        auto* state = new CApiCallbackState();
        state->py_callback = std::move(callback);
        state->callbacks.version = PJRT_API_XLA_TRANSFORM_EXTENSION_VERSION;
        state->callbacks.dtor = [](PJRT_XlaTransform_Callbacks* callbacks) {
          auto* state = reinterpret_cast<CApiCallbackState*>(
              reinterpret_cast<char*>(callbacks) -
              offsetof(CApiCallbackState, callbacks));
          delete state;
        };
        state->callbacks.transform_hlo_module = CApiTransformHloModuleCallback;

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade the plugin (TPU/GPU plugin) to a version advertising the XlaTransform extension
2. Skip registration when the extension is absent (feature-detect before registering)
3. Run transform registration only against backends known to support it

Example fix

# before
register_xla_transform_c_api(capsule, 'my_pass', 0, cb)  # CPU plugin

# after
if getattr(backend.client._plugin, 'supports_xla_transform', False):
    register_xla_transform_c_api(capsule, 'my_pass', 0, cb)

# feature-detect before registering
api = plugin._C_API
# cheap probe: registration is the only check; guard with try/except in fallback

try:
    register_xla_transform_c_api(api, name, stage, cb)
except RuntimeError as e:
    if 'XlaTransform extension' in str(e):
        logging.warning('plugin %s lacks XlaTransform; skipping transform', name)
    else:
        raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• c_api argument to register_xla_transform_c_api is not a pjrt
• c_api argument to clear_xla_transform_c_api is not a pjrt_c_
• 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/d745bc313bc7b776. Report an issue: GitHub.