{"record":{"id":"d1e7d532874e80eb","repo":"jax-ml/jax","slug":"buffer-callback-not-supported-on-platform-back","errorCode":null,"errorMessage":"`buffer_callback` not supported on {platform} backend.","messagePattern":"`buffer_callback` not supported on (.+?) backend\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/buffer_callback.py","lineNumber":246,"sourceCode":"    in_tree: Any,\n    out_tree: Any,\n    has_side_effect: bool,\n    input_output_aliases: Sequence[tuple[int, int]],\n    command_buffer_compatible: bool,\n    **_,\n):\n\n  if len(ctx.module_context.platforms) > 1:\n    raise NotImplementedError(\"multi-platform lowering for buffer_callback\")\n  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),","sourceCodeStart":228,"sourceCodeEnd":264,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/buffer_callback.py#L228-L264","documentation":"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.","triggerScenarios":"Running jax.jit code with buffer_callback on a backend not in {cpu, cuda, rocm, oneapi}, e.g. tpu, metal, or a plugin backend.","commonSituations":"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.","solutions":["Force the CPU backend for the callback-containing code: jax.config.update('jax_platforms', 'cpu') or set JAX_PLATFORMS=cpu","Replace buffer_callback with pure_callback/io_callback which have broader backend support","Guard the callback so it only executes on supported devices (lax.platform_index or jax.default_backend() check)"],"exampleFix":"// before\njax.jit(f_with_buffer_callback)(x)  # on TPU -> ValueError\n\n// after\nimport jax\nif jax.default_backend() in ('cpu', 'cuda', 'rocm'):\n    jax.jit(f_with_buffer_callback)(x)\nelse:\n    jax.jit(f_with_callback)(x)  # pure_callback variant","handlingStrategy":"validation","validationCode":"import jax\nSUPPORTED = {'cpu', 'cuda', 'rocm', 'oneapi'}\nassert jax.default_backend() in SUPPORTED, (\n    f'buffer_callback unsupported on {jax.default_backend()}')","typeGuard":"def is_supported_backend(backend: str) -> bool:\n    return backend in {'cpu', 'cuda', 'rocm', 'oneapi'}","tryCatchPattern":"try:\n    jax.jit(f_with_buffer_callback)(x)\nexcept ValueError as e:\n    if 'not supported on' in str(e) and 'buffer_callback' in str(e):\n        jax.config.update('jax_platforms', 'cpu')\n        jax.jit(f_with_buffer_callback)(x)\n    else:\n        raise","preventionTips":["Check jax.default_backend() before enabling buffer-callback debug paths","Provide a pure_callback fallback path for unsupported backends","Document backend requirements next to callback utilities"],"tags":["jax","backend","tpu","unsupported-platform","buffer-callback"],"backgroundTag":"jax-unsupported-backend-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}