jax-ml/jax · error · TypeError

No pxla_result_handler for type: {type(aval)}

Error message

No pxla_result_handler for type: {type(aval)}

What it means

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.

Source

Thrown at jax/_src/interpreters/pxla.py:283

) -> Callable[[Sequence[xc.ArrayImpl]], Any]:
  """Returns a function for handling the raw buffers of a single output aval.

  Args:
    aval: The global output AbstractValue.
    out_axis_resources: A PartitionSpec specifying the sharding of outputs.
      Used for creating GSDAs.
    global_mesh: The global device mesh that generated this output. Used
      for creating GSDAs.

  Returns:
    A function for handling the Buffers that will eventually be produced
    for this output. The function will return an object suitable for returning
    to the user, e.g. an Array.
  """
  try:
    return global_result_handlers[type(aval)](aval, out_sharding, committed)
  except KeyError as err:
    raise TypeError(
        f"No pxla_result_handler for type: {type(aval)}") from err

PxlaResultHandler = Callable[..., _jax.ResultHandler]
global_result_handlers: dict[type[core.AbstractValue], PxlaResultHandler] = {}


class InputsHandler:
  __slots__ = ("handler", "in_shardings", "in_layouts", "local_devices",
               "input_indices")

  def __init__(self, in_shardings, in_layouts, local_devices=None,
               input_indices=None):
    self.handler = partial(
        shard_args, in_shardings, in_layouts,
        [xc.ArrayCopySemantics.REUSE_INPUT] * len(in_shardings))
    self.in_shardings = in_shardings
    self.in_layouts = in_layouts
    self.local_devices = local_devices

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. If you define a custom AbstractValue, register a handler: pxla.global_result_handlers[MyAval] = handler_fn(aval, sharding, committed) -> ResultHandler
2. Update JAX — a handler missing for a built-in type is likely a fixed bug
3. Avoid producing the exotic aval as an output; convert it to a ShapedArray inside the computation

Example fix

# before
# custom aval returned from primitive with no handler

# after
from jax._src.interpreters import pxla
def _my_aval_handler(aval, out_sharding, committed):
    return lambda buf: buf
pxla.global_result_handlers[MyAval] = _my_aval_handler

from jax._src.interpreters import pxla
assert type(aval) in pxla.global_result_handlers, \
    f'register a result handler for {type(aval)}'

When it happens

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

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

Related errors

• Unsupported aval type: {type(v)}
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n

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