jax-ml/jax · error · ValueError

{primitive}.abstract_eval() method should return a tuple or

Error message

{primitive}.abstract_eval() method should return a tuple or a list iff {primitive}.multiple_results.

What it means

Every primitive's abstract_eval must return a tuple/list of abstract values if and only if primitive.multiple_results is True; a single AbstractValue otherwise. This check enforces that contract at trace time and raises when the shapes disagree (e.g. returning a bare Aval when multiple_results=True, or a 1-tuple when False).

Source

Thrown at jax/_src/interpreters/partial_eval.py:1769

                                source_info=None):
    avals = [t.aval for t in tracers]
    # TODO(mattjj): make custom_lin have hashable params.
    # TODO(dougalm): add an attribute to primitives to mark primitives with
    # effectful abstract_eval rules.
    if (primitive.ref_allocating or
        primitive.name in ("custom_lin", "call_hi_primitive_linearized",
                           "call_hi_primitive")):
      out_avals, effs = primitive.abstract_eval(*avals, **params)
    else:
      try:
        out_avals, effs = _cached_abstract_eval(primitive, *avals, **params)
      except Exception:
        # TODO(phawkins): remove this 3 months after the release of JAX v0.7.
        _verify_params_are_hashable(primitive, params)
        raise

    if isinstance(out_avals, (tuple, list)) != primitive.multiple_results:
      raise ValueError(f"{primitive}.abstract_eval() method should return "
                       f"a tuple or a list iff {primitive}.multiple_results.")
    out_avals = [out_avals] if not primitive.multiple_results else out_avals
    source_info = source_info or source_info_util.current()

    maybe_consts_out = try_constant_folding(primitive, tracers, params, out_avals)
    if maybe_consts_out is not None:
      eqn = None
      out_tracers = [self.new_const(c, source_info=source_info, aval=aval)
                     for c, aval in zip(maybe_consts_out, out_avals)]
    else:
      eqn, out_tracers = self.make_eqn(tracers, out_avals, primitive, params,
                                       effs, source_info=source_info)
    # Input-to-output tracer forwarding
    no_input_effects = not any(isinstance(e, effects.JaxprInputEffect) for e in effs)
    if eqn is not None and no_input_effects and primitive in forwarding_rules:
      in_fwd, eqn = forwarding_rules[primitive](eqn)
      for out_idx, in_idx in enumerate(in_fwd):
        if in_idx is not None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Align the flag and the return: multiple_results=True with abstract_eval returning a tuple/list of avals; multiple_results=False with a single aval
2. If your primitive genuinely has multiple outputs, set multiple_results=True when constructing core.Primitive(...) and return all avals as a tuple
3. Return exactly len(out_avals) results from the impl/rule functions too

Example fix

# before
my_prim = core.Primitive('my_prim')  # multiple_results=False
my_prim.def_abstract_eval(lambda x: (x, x))  # returns tuple

# after
my_prim = core.Primitive('my_prim', multiple_results=True)
my_prim.def_abstract_eval(lambda x: (x, x))

# in custom primitive setup
ae_result = my_prim.abstract_eval(*avals, **params)
assert isinstance(ae_result, (tuple, list)) == my_prim.multiple_results

When it happens

Trigger: Writing a custom primitive where the multiple_results flag doesn't match what abstract_eval returns. E.g. declaring multiple_results=True but abstract_eval returns a single ShapedArray, or multiple_results=False (default) but abstract_eval returns a tuple.

Common situations: New custom primitives; refactoring a single-result primitive into multi-result (or vice versa) without updating multiple_results; copy-pasted primitive scaffolding where the flag was left at its default.

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• Partitioned callback not supported with return values.
• {type(_prim).__name__} returned structured residuals from `v
• {type(_prim).__name__}.vjp_bwd should return None or a dict
• {type(_prim).__name__}.transpose should return None or a dic

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