{"record":{"id":"5617afb9424c0a87","repo":"jax-ml/jax","slug":"argument-to-register-custom-call-partitioner-was-n","errorCode":null,"errorMessage":"Argument to register_custom_call_partitioner was not a pjrt_c_api capsule.","messagePattern":"Argument to register_custom_call_partitioner was not a pjrt_c_api capsule\\.","errorType":"validation","errorClass":"absl::InvalidArgumentError","httpStatus":null,"severity":"error","filePath":"jaxlib/custom_call_sharding.cc","lineNumber":330,"sourceCode":"          sig(\n              // clang-format off\n          \"def PartiallyReplicateTiledShardingOnDims(\"\n          \"sharding: jaxlib._hlo.HloSharding, \"  // NOLINT\n          \"dims: typing.Sequence[int], /\"\n          \") -> jaxlib._hlo.HloSharding\"  // NOLINT\n                             // clang-format on\n              ));\n\n  m.def(\n      \"register_custom_call_as_batch_partitionable\",\n      [](std::string target_name, std::optional<nb::capsule> c_api) {\n        if (!c_api.has_value()) {\n          RegisterCustomCallPartitioner(\n              target_name, std::make_unique<xla::CustomCallBatchPartitioner>());\n          return;\n        }\n        if (std::string_view(c_api->name()) != \"pjrt_c_api\") {\n          throw absl::InvalidArgumentError(\n              \"Argument to register_custom_call_partitioner was not a \"\n              \"pjrt_c_api capsule.\");\n        }\n        auto* c_api_value = static_cast<const PJRT_Api*>(c_api->data());\n        PJRT_Custom_Partitioner_Extension* extension =\n            pjrt::FindExtension<PJRT_Custom_Partitioner_Extension>(\n                c_api_value,\n                PJRT_Extension_Type::PJRT_Extension_Type_Custom_Partitioner);\n        if (extension == nullptr) {\n          return;\n        }\n        PJRT_Register_Batch_Partitionable_Args args;\n        args.struct_size = PJRT_Register_Batch_Partitionable_Args_STRUCT_SIZE;\n        args.name = target_name.c_str();\n        args.name_size = target_name.size();\n        PJRT_Error* error = extension->register_batch_partitionable(&args);\n        std::unique_ptr<PJRT_Error, pjrt::PJRT_ErrorDeleter> error_ptr(\n            error, pjrt::MakeErrorDeleter(c_api_value));","sourceCodeStart":312,"sourceCodeEnd":348,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/custom_call_sharding.cc#L312-L348","documentation":"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.","triggerScenarios":"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.","commonSituations":"Plugin authors wiring custom call partitioning who pass the wrong capsule (e.g. the PJRT_Error pointer or a capsule renamed after use).","solutions":["Pass the capsule returned by jax.extend.backend.get_backend_c_api() or jaxlib.load_pjrt_plugin","Call the function without the c_api argument if you only need the default batch partitioner","Check capsule name before passing: capsule's name must equal 'pjrt_c_api'"],"exampleFix":"# before\njaxlib.register_custom_call_partitioner(name, part, c_api=wrong_capsule)\n# after\nfrom jax.extend import backend as jeb\nc_api = jeb.get_backend_c_api()\njaxlib.register_custom_call_partitioner(name, part, c_api=c_api)","handlingStrategy":"validation","validationCode":"assert c_api is None or c_api.__class__.__name__ == 'PyCapsule'","typeGuard":"def is_pjrt_c_api_capsule(cap) -> bool:\n    import ctypes\n    return cap is not None and ctypes.pythonapi.PyCapsule_GetName(cap, None) == b'pjrt_c_api'","tryCatchPattern":null,"preventionTips":["Only pass capsules from jax.extend.backend.get_backend_c_api() or load_pjrt_plugin"],"tags":["jax","xla","custom-call","capsule","validation"],"backgroundTag":"invalid-capsule-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}