{"record":{"id":"e171781a76b49a91","repo":"jax-ml/jax","slug":"host-callback-lowering-created-too-many-channels","errorCode":null,"errorMessage":"Host callback lowering created too many channels. PjRt does not support more than 65535 channels","messagePattern":"Host callback lowering created too many channels\\. PjRt does not support more than 65535 channels","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/mlir.py","lineNumber":913,"sourceCode":"        \"accessing .backend in multi-lowering setting. This can occur when \"\n        \"lowering a primitive that has not been adapted to multi-platform \"\n        \"lowering\")\n    if self.backend is not None:\n      if xb.canonicalize_platform(self.backend.platform) != self.platforms[0]:\n        if optional:\n          return None\n        raise ValueError(\n          \"the platform for the specified backend \"\n          f\"{xb.canonicalize_platform(self.backend.platform)} is different \"\n          f\"from the lowering platform {self.platforms[0]}\")\n      return self.backend\n    return xb.get_backend(self.platforms[0])\n\n  def new_channel(self) -> int:\n    channel = next(self.channel_iterator)\n    # `xla::HostCallback` requires a 16-bit channel ID.\n    if channel >= (1 << 16):\n      raise RuntimeError(\n          \"Host callback lowering created too many channels. PjRt does not\"\n          \" support more than 65535 channels\")\n    return channel\n\n  # Adds an IFRT host callback object to the context. A reference to these\n  # callbacks will be provided to IFRT during compilation so it can do things\n  # like serialize them and keep them alive.\n  def add_host_callback(self, host_callback: Any) -> None:\n    self.host_callbacks.append(host_callback)\n\n  # Keeps a value alive as long as the Python executable is alive.\n  # TODO(phawkins): this feature is problematic, because you almost certainly\n  # want to keep alive values as long as the underlying runtime executable is\n  # still alive/executing. The Python executable object may have a shorter\n  # lifetime, so it's highly likely any caller of this method is buggy.\n  def add_keepalive(self, keepalive: Any) -> None:\n    self.keepalives.append(keepalive)\n","sourceCodeStart":895,"sourceCodeEnd":931,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/mlir.py#L895-L931","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\n@jax.jit\ndef f(x):\n    for i in range(100000):\n        x = x + 1\n        jax.debug.callback(lambda v: print(v), x)  # too many channels\n\n# after\n@jax.jit\ndef body(x, _):\n    return x + 1, None\nx, _ = lax.scan(body, x, None, length=100000)\njax.debug.callback(lambda v: print(v), x)","handlingStrategy":"fallback","validationCode":"# rough static check: count callbacks that will be traced\nCB_PER_ITER = 1\nassert iterations * CB_PER_ITER < 65535, 'too many host callbacks; restructure with lax.scan'","typeGuard":null,"tryCatchPattern":"try:\n    compiled = jax.jit(f).lower(x)\nexcept RuntimeError as e:\n    if 'too many channels' in str(e):\n        f = rewrite_with_scan(f)  # emit callbacks at runtime, not trace time\n        compiled = jax.jit(f).lower(x)\n    else:\n        raise","preventionTips":["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"],"tags":["jax","host-callback","resource-limit","debugging"],"backgroundTag":"resource-limit-exceeded","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}