{"record":{"id":"cc2f6f43502939d9","repo":"jax-ml/jax","slug":"incorrect-output-shape-for-return-value-i-expe","errorCode":null,"errorMessage":"Incorrect output shape for return value #{i}: Expected: {out_aval.shape}, Actual: {out_val.shape}","messagePattern":"Incorrect output shape for return value #(.+?): Expected: (.+?), Actual: (.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":825,"sourceCode":"    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.\n      # TODO(b/238239458): fix TPU Recv to work with empty arrays.\n      non_empty_out_vals = tuple(\n          out_val\n          for out_val, result_aval in zip(out_vals, result_avals)\n          if not is_empty_shape(result_aval.shape))\n      return non_empty_out_vals\n    else:","sourceCodeStart":807,"sourceCodeEnd":843,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L807-L843","documentation":"Callback output #i has a different shape than the corresponding declared abstract value. JAX checks each returned array's shape against result_avals when the callback executes.","triggerScenarios":"A callback returning arrays whose shape differs from the shape in result_avals passed to pure_callback/io_callback (e.g. declared (3,) but returned (1,3)).","commonSituations":"Hardcoded result_shape_dtypes going stale after input shapes change; callbacks producing dynamically sized outputs; rank changes after refactoring.","solutions":["Update result_avals/result_shape_dtypes to match the actual output shapes","Make the callback reshape/pad its outputs to the declared shape before returning","Avoid dynamic output shapes; callbacks must produce the exact declared shape"],"exampleFix":"# before\nf = lambda x: x.reshape(1, -1)  # declared shape (n,)\npure_callback(f, jax.ShapeDtypeStruct(x.shape, x.dtype), x)\n# after\nf = lambda x: x\npure_callback(f, jax.ShapeDtypeStruct(x.shape, x.dtype), x)","handlingStrategy":"validation","validationCode":"def make_callback(fn, sample_in):\n    out = fn(sample_in)\n    avals = tuple(jax.ShapeDtypeStruct(np.shape(o), dtypes.canonicalize_dtype(np.dtype(o))) for o in out)\n    return avals\n# pass these avals instead of hand-written shapes","typeGuard":null,"tryCatchPattern":"try:\n    y = jax.pure_callback(fn, avals, x)\nexcept RuntimeError as e:\n    if 'Incorrect output shape' in str(e):\n        ...log and recompute avals...\n    raise","preventionTips":["Never hand-type result shapes; derive them from a dry run","Coerce outputs with np.asarray(out).reshape(aval.shape) inside the callback"],"tags":["jax","callback","shape-mismatch","runtime-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}