jax-ml/jax · error · ValueError

buffer_callback callback must not return any values.

Error message

buffer_callback callback must not return any values.

What it means

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.

Source

Thrown at jax/_src/buffer_callback.py:256

  platform = ctx.module_context.platforms[0]
  target_name = {
      "cpu": "xla_buffer_python_cpu_callback",
      "cuda": "xla_buffer_python_gpu_callback",
      "rocm": "xla_buffer_python_gpu_callback",
      "oneapi": "xla_buffer_python_gpu_callback",
  }.get(platform)
  if target_name is None:
    raise ValueError(f"`buffer_callback` not supported on {platform} backend.")

  if command_buffer_compatible and platform in ("cuda", "rocm", "oneapi"):
    target_name += "_cmd_buffer"

  def wrapped_callback(exec_ctx, *args: Any):
    args_in, args_out = util.split_list(args, [in_tree.num_leaves])
    py_args_in, py_kwargs_in = tree_util.tree_unflatten(in_tree, args_in)
    py_args_out = tree_util.tree_unflatten(out_tree, args_out)
    if callback(exec_ctx, py_args_out, *py_args_in, **py_kwargs_in) is not None:
      raise ValueError("buffer_callback callback must not return any values.")
    return ()

  ctx.module_context.add_host_callback(wrapped_callback)
  index = np.uint64(len(ctx.module_context.host_callbacks) - 1)
  rule = ffi.ffi_lowering(
      target_name,
      has_side_effect=has_side_effect,
      operand_output_aliases=dict(input_output_aliases),
  )
  return rule(ctx, *args, index=index)
mlir.register_lowering(buffer_callback_p, _buffer_callback_lowering)

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
def cb(exec_ctx, out, x):
    out[0] = x * 2
    return out  # raises

// after
def cb(exec_ctx, out, x):
    out[0] = x * 2
    return None

Defensive patterns

Strategy: validation

Validate before calling

# Validate before passing to buffer_callback
cb = my_callback
import inspect
sig = inspect.signature(cb)
# structural check: ensure no return of non-None by dry-running
out_buf = np.zeros(...)
assert cb(exec_ctx_stub, out_buf, *args) is None, 'callback must return None'

Type guard

def returns_none(cb) -> bool:
    try:
        return cb(exec_ctx_stub, out_stub, *arg_stubs) is None
    except Exception:
        return False  # structural failure, treat as unsafe

Try / catch

try:
    wrapped(...)  # runtime of the compiled callback
except ValueError as e:
    if 'must not return any values' in str(e):
        fix callback to return None and recompile
    raise

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/db63ad5207128ce3. Report an issue: GitHub.