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

c_api argument to register_xla_transform_c_api is not a pjrt

Error message

c_api argument to register_xla_transform_c_api is not a pjrt_c_api capsule.

What it means

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.

Source

Thrown at jaxlib/xla.cc:221

            break;
          case 1:
            pipeline_stage =
                xla::HloXlaTransform::PipelineStage::kPostScheduler;
            break;
          default:
            throw std::runtime_error("Invalid pipeline stage");
        }
        return xla::ClearHloXlaTransform(pipeline_stage, name);
      },
      nb::arg("name"), nb::arg("stage"));

  // 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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the plugin's pjrt_c_api capsule: typically plugin_module._C_API (named 'pjrt_c_api')
2. Verify with capsule name inspection before calling
3. Use jax.lib.xla_bridge / plugin APIs to obtain the correct capsule

Example fix

# before
register_xla_transform_c_api(jax.lib.xla_bridge.get_backend().client, ...)

# after
import jax._src.lib.xla_client as xc
capsule = jax.lib.xla_bridge.get_backend().client._plugin._C_API  # name 'pjrt_c_api'
register_xla_transform_c_api(capsule, 'pass', 0, cb)

import ctypes
# nanobind capsule; check name via its .name attribute if exposed, or plugin attribute
api = getattr(plugin_module, '_C_API', None)
if api is None:
    raise ValueError('plugin does not expose _C_API pjrt capsule')

def is_pjrt_capsule(c) -> bool:
    return type(c).__name__ == 'PyCapsule' and getattr(c, 'name', lambda: None)() == 'pjrt_c_api' if False else type(c).__name__ == 'PyCapsule'

try:
    register_xla_transform_c_api(capsule, name, stage, cb)
except ValueError as e:
    if 'pjrt_c_api capsule' in str(e):
        capsule = plugin._C_API
        register_xla_transform_c_api(capsule, name, stage, cb)
    else:
        raise

When it happens

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

Common situations: Plugin/backend integration code that grabs the wrong capsule from the PJRT plugin module; passing the module object instead of its _C_API attribute.

Related errors

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