jax-ml/jax · error · ValueError

`buffer_callback` not supported on {platform} backend.

Error message

`buffer_callback` not supported on {platform} backend.

What it means

Lowering buffer_callback maps the target platform to a registered XLA callback symbol via a dict; only 'cpu', 'cuda', 'rocm', and 'oneapi' have entries. On any other backend the lookup returns None and lowering raises ValueError naming the unsupported platform.

Source

Thrown at jax/_src/buffer_callback.py:246

    in_tree: Any,
    out_tree: Any,
    has_side_effect: bool,
    input_output_aliases: Sequence[tuple[int, int]],
    command_buffer_compatible: bool,
    **_,
):

  if len(ctx.module_context.platforms) > 1:
    raise NotImplementedError("multi-platform lowering for buffer_callback")
  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),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Force the CPU backend for the callback-containing code: jax.config.update('jax_platforms', 'cpu') or set JAX_PLATFORMS=cpu
2. Replace buffer_callback with pure_callback/io_callback which have broader backend support
3. Guard the callback so it only executes on supported devices (lax.platform_index or jax.default_backend() check)

Example fix

// before
jax.jit(f_with_buffer_callback)(x)  # on TPU -> ValueError

// after
import jax
if jax.default_backend() in ('cpu', 'cuda', 'rocm'):
    jax.jit(f_with_buffer_callback)(x)
else:
    jax.jit(f_with_callback)(x)  # pure_callback variant

import jax
SUPPORTED = {'cpu', 'cuda', 'rocm', 'oneapi'}
assert jax.default_backend() in SUPPORTED, (
    f'buffer_callback unsupported on {jax.default_backend()}')

def is_supported_backend(backend: str) -> bool:
    return backend in {'cpu', 'cuda', 'rocm', 'oneapi'}

try:
    jax.jit(f_with_buffer_callback)(x)
except ValueError as e:
    if 'not supported on' in str(e) and 'buffer_callback' in str(e):
        jax.config.update('jax_platforms', 'cpu')
        jax.jit(f_with_buffer_callback)(x)
    else:
        raise

When it happens

Trigger: Running jax.jit code with buffer_callback on a backend not in {cpu, cuda, rocm, oneapi}, e.g. tpu, metal, or a plugin backend.

Common situations: Moving a GPU debugging/snapshotting workflow to TPU; enabling a custom XLA backend or plugin; running on macOS Metal where the CPU fallback was not selected.

Related errors

• __dlpack__ device only supported for CPU, GPU and TPU pinned
• QDWH implementation is only supported on TPU
• __dlpack__ device only supported for TPU pinned host memory
• Buffer callbacks do not support transpose. Please use `jax.c
• buffer_callback callback must not return any values.

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