{"record":{"id":"b3d4a0b37c3a5166","repo":"jax-ml/jax","slug":"no-pxla-result-handler-for-type-type-aval","errorCode":null,"errorMessage":"No pxla_result_handler for type: {type(aval)}","messagePattern":"No pxla_result_handler for type: (.+?)","errorType":"error_code","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/pxla.py","lineNumber":283,"sourceCode":") -> Callable[[Sequence[xc.ArrayImpl]], Any]:\n  \"\"\"Returns a function for handling the raw buffers of a single output aval.\n\n  Args:\n    aval: The global output AbstractValue.\n    out_axis_resources: A PartitionSpec specifying the sharding of outputs.\n      Used for creating GSDAs.\n    global_mesh: The global device mesh that generated this output. Used\n      for creating GSDAs.\n\n  Returns:\n    A function for handling the Buffers that will eventually be produced\n    for this output. The function will return an object suitable for returning\n    to the user, e.g. an Array.\n  \"\"\"\n  try:\n    return global_result_handlers[type(aval)](aval, out_sharding, committed)\n  except KeyError as err:\n    raise TypeError(\n        f\"No pxla_result_handler for type: {type(aval)}\") from err\n\nPxlaResultHandler = Callable[..., _jax.ResultHandler]\nglobal_result_handlers: dict[type[core.AbstractValue], PxlaResultHandler] = {}\n\n\nclass InputsHandler:\n  __slots__ = (\"handler\", \"in_shardings\", \"in_layouts\", \"local_devices\",\n               \"input_indices\")\n\n  def __init__(self, in_shardings, in_layouts, local_devices=None,\n               input_indices=None):\n    self.handler = partial(\n        shard_args, in_shardings, in_layouts,\n        [xc.ArrayCopySemantics.REUSE_INPUT] * len(in_shardings))\n    self.in_shardings = in_shardings\n    self.in_layouts = in_layouts\n    self.local_devices = local_devices","sourceCodeStart":265,"sourceCodeEnd":301,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/pxla.py#L265-L301","documentation":"When a compiled computation returns, pxla looks up a result handler for each output's abstract value type in global_result_handlers (registered for ShapedArray, token, etc.). An unregistered aval type means JAX knows how to produce the value internally but not how to wrap it for return to the user, so it raises TypeError.","triggerScenarios":"A primitive's abstract_eval returns a non-standard AbstractValue (a custom aval class) without registering a corresponding entry in pxla.global_result_handlers; or an internal aval type that lacks a handler reaches execution.","commonSituations":"Custom avals from extensions not registering result handlers; JAX version mismatches where a new aval type exists but its handler registration was missed (often an internal bug).","solutions":["If you define a custom AbstractValue, register a handler: pxla.global_result_handlers[MyAval] = handler_fn(aval, sharding, committed) -> ResultHandler","Update JAX — a handler missing for a built-in type is likely a fixed bug","Avoid producing the exotic aval as an output; convert it to a ShapedArray inside the computation"],"exampleFix":"# before\n# custom aval returned from primitive with no handler\n\n# after\nfrom jax._src.interpreters import pxla\ndef _my_aval_handler(aval, out_sharding, committed):\n    return lambda buf: buf\npxla.global_result_handlers[MyAval] = _my_aval_handler","handlingStrategy":"validation","validationCode":"from jax._src.interpreters import pxla\nassert type(aval) in pxla.global_result_handlers, \\\n    f'register a result handler for {type(aval)}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Register pxla.global_result_handlers entries for custom avals in your extension's init","Keep custom avals internal, convert to ShapedArray at API boundaries"],"tags":["jax","pxla","custom-aval","result-handler","extension-api"],"backgroundTag":"missing-type-handler-registration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}