{"record":{"id":"c549c6c25e644e89","repo":"jax-ml/jax","slug":"primitive-abstract-eval-method-should-return-a","errorCode":null,"errorMessage":"{primitive}.abstract_eval() method should return a tuple or a list iff {primitive}.multiple_results.","messagePattern":"(.+?)\\.abstract_eval\\(\\) method should return a tuple or a list iff (.+?)\\.multiple_results\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/partial_eval.py","lineNumber":1769,"sourceCode":"                                source_info=None):\n    avals = [t.aval for t in tracers]\n    # TODO(mattjj): make custom_lin have hashable params.\n    # TODO(dougalm): add an attribute to primitives to mark primitives with\n    # effectful abstract_eval rules.\n    if (primitive.ref_allocating or\n        primitive.name in (\"custom_lin\", \"call_hi_primitive_linearized\",\n                           \"call_hi_primitive\")):\n      out_avals, effs = primitive.abstract_eval(*avals, **params)\n    else:\n      try:\n        out_avals, effs = _cached_abstract_eval(primitive, *avals, **params)\n      except Exception:\n        # TODO(phawkins): remove this 3 months after the release of JAX v0.7.\n        _verify_params_are_hashable(primitive, params)\n        raise\n\n    if isinstance(out_avals, (tuple, list)) != primitive.multiple_results:\n      raise ValueError(f\"{primitive}.abstract_eval() method should return \"\n                       f\"a tuple or a list iff {primitive}.multiple_results.\")\n    out_avals = [out_avals] if not primitive.multiple_results else out_avals\n    source_info = source_info or source_info_util.current()\n\n    maybe_consts_out = try_constant_folding(primitive, tracers, params, out_avals)\n    if maybe_consts_out is not None:\n      eqn = None\n      out_tracers = [self.new_const(c, source_info=source_info, aval=aval)\n                     for c, aval in zip(maybe_consts_out, out_avals)]\n    else:\n      eqn, out_tracers = self.make_eqn(tracers, out_avals, primitive, params,\n                                       effs, source_info=source_info)\n    # Input-to-output tracer forwarding\n    no_input_effects = not any(isinstance(e, effects.JaxprInputEffect) for e in effs)\n    if eqn is not None and no_input_effects and primitive in forwarding_rules:\n      in_fwd, eqn = forwarding_rules[primitive](eqn)\n      for out_idx, in_idx in enumerate(in_fwd):\n        if in_idx is not None:","sourceCodeStart":1751,"sourceCodeEnd":1787,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/partial_eval.py#L1751-L1787","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","If your primitive genuinely has multiple outputs, set multiple_results=True when constructing core.Primitive(...) and return all avals as a tuple","Return exactly len(out_avals) results from the impl/rule functions too"],"exampleFix":"# before\nmy_prim = core.Primitive('my_prim')  # multiple_results=False\nmy_prim.def_abstract_eval(lambda x: (x, x))  # returns tuple\n\n# after\nmy_prim = core.Primitive('my_prim', multiple_results=True)\nmy_prim.def_abstract_eval(lambda x: (x, x))","handlingStrategy":"validation","validationCode":"# in custom primitive setup\nae_result = my_prim.abstract_eval(*avals, **params)\nassert isinstance(ae_result, (tuple, list)) == my_prim.multiple_results","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Set multiple_results at Primitive construction to match abstract_eval's return shape","Add a smoke test that binds and evaluates each custom primitive"],"tags":["jax","custom-primitives","abstract-eval","multiple-results","api-contract"],"backgroundTag":"api-contract-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}