jax-ml/jax · error · NotImplementedError

Ordered debug_print is not supported on Pallas.

Error message

Ordered debug_print is not supported on Pallas.

What it means

jax.debug.print(ordered=True) is not supported inside Pallas kernels on GPU. Ordered printing requires a deterministic callback ordering across devices/replicas which the Pallas lowering cannot provide.

Source

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

@register_lowering_rule(
    debugging.debug_print_p, mgpu.LoweringSemantics.Warpgroup
)
@register_lowering_rule(debugging.debug_print_p, *gpu_core.WGxWARP_SEMANTICS)
def _debug_print_lowering_rule(
    ctx: LoweringRuleContext,
    *args,
    fmt,
    ordered,
    partitioned,
    in_tree,
    static_args,
    np_printoptions,
    has_placeholders,
    logging_record,
):
  del partitioned, np_printoptions, has_placeholders
  if ordered:
    raise NotImplementedError("Ordered debug_print is not supported on Pallas.")
  args, kwargs = debugging.merge_callback_args(in_tree, args, static_args)
  if kwargs:
    raise ValueError(
        "Only positional arguments are supported by debug_print on Pallas."
    )
  primitives.check_debug_print_format(fmt, *args)
  if not any(aval.shape for aval in ctx.avals_in):
    scope = mgpu.ThreadSubset.WARPGROUP
    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
      scope = mgpu.ThreadSubset.WARP
    mgpu.debug_print(
        fmt,
        *(
            _ensure_ir_value(arg, aval.dtype)
            for arg, aval in zip(args, ctx.avals_in)
        ),
        scope=scope
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop ordered=True and accept interleaved output
2. Use jax.debug.print without ordering, or print summary scalars after the kernel returns
3. For deterministic debugging, write values to an output buffer and inspect after launch

Example fix

// before
jax.debug.print("x={x}", x=x, ordered=True)
// after
jax.debug.print("x={x}", x=x)

if ordered:
    ordered = False  # not supported under Pallas

When it happens

Trigger: Calling jax.debug.print(..., ordered=True) inside a Pallas GPU kernel function.

Common situations: Copy-pasting debug prints from JIT-compiled JAX code (where ordered=True is common to keep output readable) into a Pallas kernel.

Related errors

• Only positional arguments are supported by debug_print on Pa
• Can only print scalars in warp-level code.
• debug_print only supports printing of scalar values, or a si
• 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/6475b7a976d68eaf. Report an issue: GitHub.