jax-ml/jax · error · ValueError

`buffer_callback` not supported on

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.

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…
• QDWH implementation is only supported on TPU
• Acc ref dtype must be float32 or int32, got
• Acc ref must be at least 2D, got shape
• Accumulators are not available on TPU

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

Appendix: 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)