jax-ml/jax · error · absl::InvalidArgumentError

Argument to register_custom_call_partitioner was not a pjrt_

Error message

Argument to register_custom_call_partitioner was not a pjrt_c_api capsule.

What it means

register_custom_call_partitioner optionally takes a pjrt_c_api capsule to attach a C-API custom partitioner extension. If the passed capsule's name is anything other than 'pjrt_c_api', the argument is rejected as invalid.

Source

Thrown at jaxlib/custom_call_sharding.cc:330

          sig(
              // clang-format off
          "def PartiallyReplicateTiledShardingOnDims("
          "sharding: jaxlib._hlo.HloSharding, "  // NOLINT
          "dims: typing.Sequence[int], /"
          ") -> jaxlib._hlo.HloSharding"  // NOLINT
                             // clang-format on
              ));

  m.def(
      "register_custom_call_as_batch_partitionable",
      [](std::string target_name, std::optional<nb::capsule> c_api) {
        if (!c_api.has_value()) {
          RegisterCustomCallPartitioner(
              target_name, std::make_unique<xla::CustomCallBatchPartitioner>());
          return;
        }
        if (std::string_view(c_api->name()) != "pjrt_c_api") {
          throw absl::InvalidArgumentError(
              "Argument to register_custom_call_partitioner was not a "
              "pjrt_c_api capsule.");
        }
        auto* c_api_value = static_cast<const PJRT_Api*>(c_api->data());
        PJRT_Custom_Partitioner_Extension* extension =
            pjrt::FindExtension<PJRT_Custom_Partitioner_Extension>(
                c_api_value,
                PJRT_Extension_Type::PJRT_Extension_Type_Custom_Partitioner);
        if (extension == nullptr) {
          return;
        }
        PJRT_Register_Batch_Partitionable_Args args;
        args.struct_size = PJRT_Register_Batch_Partitionable_Args_STRUCT_SIZE;
        args.name = target_name.c_str();
        args.name_size = target_name.size();
        PJRT_Error* error = extension->register_batch_partitionable(&args);
        std::unique_ptr<PJRT_Error, pjrt::PJRT_ErrorDeleter> error_ptr(
            error, pjrt::MakeErrorDeleter(c_api_value));

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the capsule returned by jax.extend.backend.get_backend_c_api() or jaxlib.load_pjrt_plugin
2. Call the function without the c_api argument if you only need the default batch partitioner
3. Check capsule name before passing: capsule's name must equal 'pjrt_c_api'

Example fix

# before
jaxlib.register_custom_call_partitioner(name, part, c_api=wrong_capsule)
# after
from jax.extend import backend as jeb
c_api = jeb.get_backend_c_api()
jaxlib.register_custom_call_partitioner(name, part, c_api=c_api)

assert c_api is None or c_api.__class__.__name__ == 'PyCapsule'

def is_pjrt_c_api_capsule(cap) -> bool:
    import ctypes
    return cap is not None and ctypes.pythonapi.PyCapsule_GetName(cap, None) == b'pjrt_c_api'

When it happens

Trigger: Calling jaxlib.register_custom_call_partitioner(name, partitioner, c_api) with a capsule obtained from something other than jaxlib.load_pjrt_plugin / jax.extend.backend.get_backend_c_api(), or passing a manually-constructed capsule.

Common situations: Plugin authors wiring custom call partitioning who pass the wrong capsule (e.g. the PJRT_Error pointer or a capsule renamed after use).

Related errors

• c_api argument to load_pjrt_plugin is not a pjrt_c_api capsu
• Argument to get_c_api_topology was not a pjrt_c_api capsule.
• Invalid pipeline stage
• {full_name} must be a pytree prefix with bool leaves or a tu
• multi-platform lowering for buffer_callback

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