jax-ml/jax · error · RuntimeError

Host callback lowering created too many channels. PjRt does

Error message

Host callback lowering created too many channels. PjRt does not support more than 65535 channels

What it means

Host callbacks (e.g. from jax.debug.callback / experimental host_callback) communicate with the runtime over a 16-bit channel ID, so at most 65535 channels can exist per lowering. Each new callback in a single compilation consumes a channel; exceeding the limit raises this RuntimeError.

Source

Thrown at jax/_src/interpreters/mlir.py:913

        "accessing .backend in multi-lowering setting. This can occur when "
        "lowering a primitive that has not been adapted to multi-platform "
        "lowering")
    if self.backend is not None:
      if xb.canonicalize_platform(self.backend.platform) != self.platforms[0]:
        if optional:
          return None
        raise ValueError(
          "the platform for the specified backend "
          f"{xb.canonicalize_platform(self.backend.platform)} is different "
          f"from the lowering platform {self.platforms[0]}")
      return self.backend
    return xb.get_backend(self.platforms[0])

  def new_channel(self) -> int:
    channel = next(self.channel_iterator)
    # `xla::HostCallback` requires a 16-bit channel ID.
    if channel >= (1 << 16):
      raise RuntimeError(
          "Host callback lowering created too many channels. PjRt does not"
          " support more than 65535 channels")
    return channel

  # Adds an IFRT host callback object to the context. A reference to these
  # callbacks will be provided to IFRT during compilation so it can do things
  # like serialize them and keep them alive.
  def add_host_callback(self, host_callback: Any) -> None:
    self.host_callbacks.append(host_callback)

  # Keeps a value alive as long as the Python executable is alive.
  # TODO(phawkins): this feature is problematic, because you almost certainly
  # want to keep alive values as long as the underlying runtime executable is
  # still alive/executing. The Python executable object may have a shorter
  # lifetime, so it's highly likely any caller of this method is buggy.
  def add_keepalive(self, keepalive: Any) -> None:
    self.keepalives.append(keepalive)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move the callback out of the unrolled loop: call it once on the collected results, or use vmap/scan (lax.scan) so the callback executes per-iteration at runtime, not per-unroll
Batch the data you want to send to the host and emit a single callback
Reduce instrumentation density during debugging

Example fix

# before
@jax.jit
def f(x):
    for i in range(100000):
        x = x + 1
        jax.debug.callback(lambda v: print(v), x)  # too many channels

# after
@jax.jit
def body(x, _):
    return x + 1, None
x, _ = lax.scan(body, x, None, length=100000)
jax.debug.callback(lambda v: print(v), x)

Defensive patterns

Strategy: fallback

Validate before calling

# rough static check: count callbacks that will be traced
CB_PER_ITER = 1
assert iterations * CB_PER_ITER < 65535, 'too many host callbacks; restructure with lax.scan'

Try / catch

try:
    compiled = jax.jit(f).lower(x)
except RuntimeError as e:
    if 'too many channels' in str(e):
        f = rewrite_with_scan(f)  # emit callbacks at runtime, not trace time
        compiled = jax.jit(f).lower(x)
    else:
        raise

Prevention

Never put jax.debug.callback/print inside Python for-loops under jit
Use lax.scan/vmap so one callback executes per iteration at runtime
Batch debug output and emit a single callback

When it happens

Trigger: A single compiled computation containing more than 65535 host callbacks — typically loops unrolled at trace time that each emit a jax.debug.callback, host_callback.id_tap, or print inside a python loop run under jit.

Common situations: Using jax.debug.print/callback inside a Python for-loop that Python-unrolls under jit; large batched tap instrumentation; migrating from the deprecated host_callback package which emits many taps; unrolled loops over tens of thousands of iterations.

Related errors

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