{"record":{"id":"24ca3519a73878c4","repo":"jax-ml/jax","slug":"multi-platform-lowering-for-buffer-callback","errorCode":null,"errorMessage":"multi-platform lowering for buffer_callback","messagePattern":"multi-platform lowering for buffer_callback","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/buffer_callback.py","lineNumber":237,"sourceCode":"batching.primitive_batchers[buffer_callback_p] = functools.partial(\n    ffi.ffi_batching_rule, buffer_callback_p\n)\n\n\ndef _buffer_callback_lowering(\n    ctx: mlir.LoweringRuleContext,\n    *args: Any,\n    callback,\n    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:","sourceCodeStart":219,"sourceCodeEnd":255,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/buffer_callback.py#L219-L255","documentation":"_buffer_callback_lowering compiles buffer_callback to a platform-specific XLA external custom call target, and it only knows how to pick a single target name. When the compilation context targets more than one platform simultaneously (multi-platform lowering), there is no single backend-specific callback symbol to emit, so lowering fails with NotImplementedError.","triggerScenarios":"Compiling/jitting a function containing buffer_callback when ctx.module_context.platforms has length > 1, e.g. exporting a multi-platform artifact (jax.export / jaxlib MLIR export targeting cpu+cuda) that includes a buffer callback.","commonSituations":"Using jax.experimental.export or multi-platform AOT compilation with debugger/buffer-inspection callbacks still in the program; migrating single-platform pipelines to multi-platform export without removing host callbacks.","solutions":["Restrict compilation to a single platform (drop the extra targets from the export/compile platform list)","Remove or gate the buffer_callback before multi-platform export (strip debugging callbacks)","Export the callback-free model and attach host-side logic per platform after loading"],"exampleFix":"// before\nexp = jax.experimental.export.export(\n    jax.jit(f), platforms=('cpu', 'cuda'))(x)  # f contains buffer_callback\n\n// after\nexp = jax.experimental.export.export(\n    jax.jit(f_without_callback), platforms=('cpu', 'cuda'))(x)","handlingStrategy":"validation","validationCode":"from jax.experimental.export import export\nplatforms = ('cpu',)  # keep single-platform when buffer_callback is present\nassert len(platforms) == 1, 'multi-platform export cannot include buffer_callback'","typeGuard":null,"tryCatchPattern":"try:\n    exported = export(jax.jit(f), platforms=platforms)(x)\nexcept NotImplementedError as e:\n    if 'multi-platform lowering for buffer_callback' in str(e):\n        platforms = platforms[:1]\n        exported = export(jax.jit(f), platforms=platforms)(x)\n    else:\n        raise","preventionTips":["Strip debugging/buffer callbacks before AOT export","Keep a callback-free variant of the model for export pipelines","Pin export platform lists in one config location so they are easy to audit"],"tags":["jax","xla","multi-platform","lowering","export"],"backgroundTag":"xla-multi-platform-lowering-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}