jax-ml/jax · error · ValueError

Only positional arguments are supported by debug_print on Pa

Error message

Only positional arguments are supported by debug_print on Pallas.

What it means

Pallas GPU debug_print lowering only handles positional arguments; passing keyword arguments (fmt kwargs for values) raises ValueError before codegen.

Source

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

@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
    )
  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.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass values positionally: jax.debug.print("{}", value)
Use {}-style positional placeholders matching the argument order
Or move the print outside the kernel

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

# only positional: jax.debug.print('{} {}', a, b)

Prevention

Use positional placeholders in Pallas debug prints
Lint for debug_print kwargs in kernel modules

When it happens

Trigger: jax.debug.print("v={v}", v=value) inside a Pallas kernel — the kwarg form — triggers the error even though it's idiomatic in plain JAX.

Common situations: Using the named-placeholder style of debug_print (which formats by keyword) inside kernels; it works under jit but not under Pallas.

Related errors

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