{"record":{"id":"8a304c7ae4e21691","repo":"jax-ml/jax","slug":"partitioned-callback-not-supported-with-return-val","errorCode":null,"errorMessage":"Partitioned callback not supported with return values.","messagePattern":"Partitioned callback not supported with return values\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":808,"sourceCode":"      False, then `callback` is called on all shards.\n    sharding: The sharding of the callback.\n\n  Returns:\n    A tuple of MLIR result values, a new token (if any), and the host callback\n    object.\n  \"\"\"\n  if len(ctx.module_context.platforms) > 1:\n    raise NotImplementedError(\"multi-platform lowering for python_callback\")\n  platform = ctx.module_context.platforms[0]\n  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}: \"","sourceCodeStart":790,"sourceCodeEnd":826,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L790-L826","documentation":"A partitioned callback was requested but the callback declares result values (result_avals). Partitioned callbacks must be side-effecting procedures with no return values.","triggerScenarios":"Calling io_callback/pure_callback with _partitioned=True while passing a non-empty result_avals tuple, i.e. the callback returns values.","commonSituations":"Adapting existing pure_callback code (which returns values) to run partitioned; misunderstanding that partitioned callbacks are emit-only.","solutions":["Remove return values from the callback and return nothing (make it a side-effecting callback)","If you need results, use a non-partitioned pure_callback/io_callback"],"exampleFix":"# before\nio_callback(lambda x: x * 2, make_jaxpr(...).out_avals, x, _partitioned=True)\n# after\nio_callback(lambda x: None, (), x, _partitioned=True)  # side effects only","handlingStrategy":"validation","validationCode":"def make_partitioned_callback(fn, result_avals):\n    if result_avals:\n        raise ValueError('use non-partitioned callback for value-returning fns')\n    return partial(io_callback, fn, result_avals, _partitioned=True)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep partitioned callbacks side-effect-only by convention","Assert result_avals == () in your wrapper helpers"],"tags":["jax","callback","sharding","api-contract"],"backgroundTag":"invalid-argument-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}