{"record":{"id":"1b6e2a7fbdd7a9a8","repo":"jax-ml/jax","slug":"mismatched-number-of-outputs-from-callback-expect","errorCode":null,"errorMessage":"Mismatched number of outputs from callback. Expected: {}, Actual: {}","messagePattern":"Mismatched number of outputs from callback\\. Expected: (.+?), Actual: (.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":818,"sourceCode":"  if platform not in {\"cpu\", \"cuda\", \"rocm\", \"tpu\", \"oneapi\"}:\n    raise ValueError(\n        f\"`EmitPythonCallback` not supported on {platform} backend.\")\n  if partitioned:\n    if platform not in {\"cpu\", \"cuda\", \"rocm\", \"oneapi\"}:\n      raise NotImplementedError(\n          f\"Partitioned callback not implemented on {platform} backend.\")\n    if result_avals:\n      raise ValueError(\"Partitioned callback not supported with return values.\")\n  backend: xc.Client = cast(xc.Client, ctx.module_context.get_backend())\n  result_shapes = [_aval_to_xla_shape(aval) for aval in result_avals]\n  operand_shapes = [_aval_to_xla_shape(aval) for aval in operand_avals]\n\n  # First we apply checks to ensure output shapes and dtypes match the expected\n  # ones.\n  def _wrapped_callback(*args):\n    out_vals = callback(*args)\n    if len(out_vals) != len(result_avals):\n      raise RuntimeError(\n          \"Mismatched number of outputs from callback. \"\n          \"Expected: {}, Actual: {}\".format(len(result_avals), len(out_vals)))\n    # Handle Python literals, and custom arrays, e.g., tf.Tensor.\n    out_vals = tuple(dtypes.canonicalize_value(np.asarray(a)) for a in out_vals)\n    for i, (out_val, out_aval) in enumerate(zip(out_vals, result_avals)):\n      if out_val.shape != out_aval.shape:\n        raise RuntimeError(\n            f\"Incorrect output shape for return value #{i}: \"\n            f\"Expected: {out_aval.shape}, Actual: {out_val.shape}\")\n      if out_val.dtype != out_aval.dtype:\n        raise RuntimeError(\n            f\"Incorrect output dtype for return value #{i}: \"\n            f\"Expected: {out_aval.dtype}, Actual: {out_val.dtype}\")\n\n    if platform == \"tpu\":\n      # On TPU we cannot receive empty arrays. So, we return from the wrapped\n      # callback only the non-empty results, and we will create empty constants\n      # in the receiving computation.","sourceCodeStart":800,"sourceCodeEnd":836,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L800-L836","documentation":"The wrapped Python callback returned a different number of outputs than declared in result_avals. JAX validates callback results at runtime against the promised abstract values.","triggerScenarios":"A pure_callback/io_callback whose wrapped function returns a tuple of length different from len(result_avals); e.g. returning a single array when two were declared, or returning None.","commonSituations":"Callback returns a scalar instead of a 1-tuple; callback branches and returns different tuple sizes; result_shape/dtype declaration out of sync with the function after refactoring.","solutions":["Make the callback return exactly as many outputs as declared in result_avals (wrap single values in a tuple)","Re-check the result_shape_dtypes declaration matches the function's actual return","Return () explicitly for zero-result callbacks instead of None"],"exampleFix":"# before\nf = lambda x: (x + 1)  # declared 2 results\npure_callback(f, (shapes, shapes2), x)\n# after\nf = lambda x: (x + 1, x + 2)\npure_callback(f, (shapes, shapes2), x)","handlingStrategy":"validation","validationCode":"def checked_pure_callback(fn, result_avals, *args):\n    out = fn(*args)  # dry-run on sample inputs in tests\n    assert len(jtu.tree_leaves(out)) == len(result_avals), (len(out), len(result_avals))\n    return jax.pure_callback(fn, result_avals, *args)","typeGuard":null,"tryCatchPattern":"try:\n    y = jax.pure_callback(fn, avals, x)\nexcept RuntimeError as e:\n    if 'Mismatched number of outputs' in str(e):\n        raise ValueError('callback arity drift; re-sync result_avals') from e\n    raise","preventionTips":["Unit-test callbacks on concrete inputs before jit","Single source of truth for result types: derive avals from a sample output"],"tags":["jax","callback","runtime-validation","shape-mismatch"],"backgroundTag":"callback-return-value-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}