jax-ml/jax · error · NotImplementedError

multi-platform lowering for buffer_callback

Error message

multi-platform lowering for buffer_callback

What it means

_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.

Source

Thrown at jax/_src/buffer_callback.py:237

batching.primitive_batchers[buffer_callback_p] = functools.partial(
    ffi.ffi_batching_rule, buffer_callback_p
)


def _buffer_callback_lowering(
    ctx: mlir.LoweringRuleContext,
    *args: Any,
    callback,
    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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restrict compilation to a single platform (drop the extra targets from the export/compile platform list)
2. Remove or gate the buffer_callback before multi-platform export (strip debugging callbacks)
3. Export the callback-free model and attach host-side logic per platform after loading

Example fix

// before
exp = jax.experimental.export.export(
    jax.jit(f), platforms=('cpu', 'cuda'))(x)  # f contains buffer_callback

// after
exp = jax.experimental.export.export(
    jax.jit(f_without_callback), platforms=('cpu', 'cuda'))(x)

from jax.experimental.export import export
platforms = ('cpu',)  # keep single-platform when buffer_callback is present
assert len(platforms) == 1, 'multi-platform export cannot include buffer_callback'

try:
    exported = export(jax.jit(f), platforms=platforms)(x)
except NotImplementedError as e:
    if 'multi-platform lowering for buffer_callback' in str(e):
        platforms = platforms[:1]
        exported = export(jax.jit(f), platforms=platforms)(x)
    else:
        raise

When it happens

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

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

Related errors

• accessing .backend in multi-lowering setting. This can occur
• Nesting `compute_on` with different compute types is not all
• No VJP is available
• Duplicate serialization registration for type `{nodetype}`.
• Duplicate serialization registration for serialized_name `{s

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