jax-ml/jax · error · NotImplementedError

debug_print only supports printing of scalar values, or a si

Error message

debug_print only supports printing of scalar values, or a single array value when using the Mosaic GPU backend.

What it means

Mosaic GPU's debug_print supports (a) any number of scalar arguments, or (b) exactly one array argument. Passing multiple arrays (or a mix where more than one input has shape) raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3549

    mgpu.debug_print(
        fmt,
        *(
            _ensure_ir_value(arg, aval.dtype)
            for arg, aval in zip(args, ctx.avals_in)
        ),
        scope=scope
    )
  elif len(ctx.avals_in) == 1:
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      raise NotImplementedError("Can only print scalars in warp-level code.")
    [arg] = args
    if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
      mgpu.dialect.debug_print(fmt, arg)
    else:
      arg.debug_print(fmt)

  else:
    raise NotImplementedError(
        "debug_print only supports printing of scalar values, or a single array"
        " value when using the Mosaic GPU backend."
    )

  return ()


@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Warpgroup)
def _run_scoped_lowering_rule(
    ctx: LoweringRuleContext,
    *consts,
    jaxpr: jax_core.Jaxpr,
    collective_axes,
    **_,
):
  if pallas_core.poison_buffers_enabled():
    raise NotImplementedError("Buffer poisoning is not supported on GPU yet.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split into separate debug_print calls, one per array
2. Or print scalars (e.g. per-tile sums) for each array in a single call
3. Remember the rule: many scalars OK, one array OK, multiple arrays not

Example fix

// before
jax.debug.print("{} {}", a, b)
// after
jax.debug.print("{}", a)
jax.debug_print("{}", b)

arrays = [a for a in args if a.shape]
assert len(arrays) <= 1, 'one array max per debug_print'

def debug_print_ok(avals) -> bool:
    return sum(1 for a in avals if a.shape) <= 1

When it happens

Trigger: jax.debug.print("{} {}", arr1, arr2) inside a Pallas GPU kernel where both args are arrays with non-empty shapes.

Common situations: Trying to print two tiles at once during kernel debugging; combining scalars and arrays in one call in the array branch.

Related errors

• Ordered debug_print is not supported on Pallas.
• Only positional arguments are supported by debug_print on Pa
• Can only print scalars in warp-level code.
• The format string expects {n_placeholders} argument{'' if n_
• Only positional arguments are supported by debug_print on Pa

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