{"record":{"id":"db63ad5207128ce3","repo":"jax-ml/jax","slug":"buffer-callback-callback-must-not-return-any-value","errorCode":null,"errorMessage":"buffer_callback callback must not return any values.","messagePattern":"buffer_callback callback must not return any values\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/buffer_callback.py","lineNumber":256,"sourceCode":"  platform = ctx.module_context.platforms[0]\n  target_name = {\n      \"cpu\": \"xla_buffer_python_cpu_callback\",\n      \"cuda\": \"xla_buffer_python_gpu_callback\",\n      \"rocm\": \"xla_buffer_python_gpu_callback\",\n      \"oneapi\": \"xla_buffer_python_gpu_callback\",\n  }.get(platform)\n  if target_name is None:\n    raise ValueError(f\"`buffer_callback` not supported on {platform} backend.\")\n\n  if command_buffer_compatible and platform in (\"cuda\", \"rocm\", \"oneapi\"):\n    target_name += \"_cmd_buffer\"\n\n  def wrapped_callback(exec_ctx, *args: Any):\n    args_in, args_out = util.split_list(args, [in_tree.num_leaves])\n    py_args_in, py_kwargs_in = tree_util.tree_unflatten(in_tree, args_in)\n    py_args_out = tree_util.tree_unflatten(out_tree, args_out)\n    if callback(exec_ctx, py_args_out, *py_args_in, **py_kwargs_in) is not None:\n      raise ValueError(\"buffer_callback callback must not return any values.\")\n    return ()\n\n  ctx.module_context.add_host_callback(wrapped_callback)\n  index = np.uint64(len(ctx.module_context.host_callbacks) - 1)\n  rule = ffi.ffi_lowering(\n      target_name,\n      has_side_effect=has_side_effect,\n      operand_output_aliases=dict(input_output_aliases),\n  )\n  return rule(ctx, *args, index=index)\nmlir.register_lowering(buffer_callback_p, _buffer_callback_lowering)\n","sourceCodeStart":238,"sourceCodeEnd":268,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/buffer_callback.py#L238-L268","documentation":"At lowering time buffer_callback wraps the user callback and executes it with (exec_ctx, outputs, *args, **kwargs); the wrapper asserts the callback's Python return value is None because the callback communicates results by mutating the output buffers in place, not by returning them. Any non-None return (including accidentally returning a value, a tuple, or a truthy sentinel) triggers this ValueError from the wrapped host callback at runtime.","triggerScenarios":"Passing a callback to buffer_callback that has a return statement returning anything non-None, e.g. `def cb(ctx, out, x): return out[0] += 1` or a function whose last expression evaluates to a value.","commonSituations":"Refactoring a pure_callback (which must return results) into a buffer_callback without removing the return; writing `return None` vs implicit-return confusion; callbacks written as lambdas that evaluate to a value.","solutions":["Make the callback return None: mutate output buffers in place and end with bare return / no return","If you meant to return values, use jax.pure_callback instead of buffer_callback","Audit the callback for implicit returns (last expression) and lambda bodies"],"exampleFix":"// before\ndef cb(exec_ctx, out, x):\n    out[0] = x * 2\n    return out  # raises\n\n// after\ndef cb(exec_ctx, out, x):\n    out[0] = x * 2\n    return None","handlingStrategy":"validation","validationCode":"# Validate before passing to buffer_callback\ncb = my_callback\nimport inspect\nsig = inspect.signature(cb)\n# structural check: ensure no return of non-None by dry-running\nout_buf = np.zeros(...)\nassert cb(exec_ctx_stub, out_buf, *args) is None, 'callback must return None'","typeGuard":"def returns_none(cb) -> bool:\n    try:\n        return cb(exec_ctx_stub, out_stub, *arg_stubs) is None\n    except Exception:\n        return False  # structural failure, treat as unsafe","tryCatchPattern":"try:\n    wrapped(...)  # runtime of the compiled callback\nexcept ValueError as e:\n    if 'must not return any values' in str(e):\n        fix callback to return None and recompile\n    raise","preventionTips":["End buffer callbacks with bare return or return None explicitly","Never convert pure_callback bodies to buffer_callback without deleting returns","Lint callbacks with a unit test asserting the return value is None"],"tags":["jax","callback","host-callback","return-value","buffer-callback"],"backgroundTag":"callback-return-value-contract","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}