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

Raised by the Pallas→Triton lowering rule for jax.debug_print: after merging callback args, any remaining keyword arguments are rejected because the Triton backend only implements printing with positional values. The check mirrors a similar restriction against placeholder formatting. It is thrown at compile/lowering time, before the kernel runs on GPU.

Source

Thrown at jax/_src/pallas/triton/lowering.py:1332

    fmt: str,
    ordered,
    partitioned,
    in_tree,
    static_args,
    np_printoptions,
    has_placeholders,
    logging_record,
):
  del partitioned, np_printoptions
  if ordered:
    raise NotImplementedError("Ordered debug_print is not supported on Pallas.")
  if has_placeholders:
    raise ValueError(
        "pl.debug_print() does not support placeholders when lowering to Triton"
    )
  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."
    )

  tt_dialect.print_(
      f" {fmt} ",
      hex=False,
      args=args,
      is_signed=ir.DenseI32ArrayAttr.get([
          jnp.issubdtype(aval.dtype, jnp.signedinteger) for aval in ctx.avals_in
      ]),
  )
  return ()


def _set_attr(v: ir.Value, name: str, attr: ir.Attribute) -> None:
  if not isinstance(v, ir.BlockArgument):
    v.owner.attributes[name] = attr  # pyrefly: ignore[missing-attribute]
    return

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass all values to pl.debug_print positionally, e.g. pl.debug_print('x={} y={}', x, y) instead of pl.debug_print('x={x}', x=x)
Remove any placeholder-style format strings as well; pre-format into a simple '{}/{}/...' positional template
If you need named debugging, compute a tuple and print it positionally or fall back to jax.debug.print outside the kernel

Example fix

// before
pl.debug_print("a={a} b={b}", a=block, b=acc)
// after
pl.debug_print("a={} b={}", block, acc)

Defensive patterns

Strategy: validation

Validate before calling

def check_debug_print_args(fmt, args, kwargs):
    if kwargs:
        raise TypeError("pl.debug_print on Triton only supports positional args")
    if '{' in fmt and not re.fullmatch(r'(\{\})*', fmt.replace(' ', '')):
        raise TypeError("use positional {} placeholders, not named placeholders")
    return True

Prevention

Always call pl.debug_print(fmt, *values) with positional args and '{}' placeholders
Wrap debug printing in a helper that enforces positional args so all call sites stay Triton-safe

When it happens

Trigger: Calling pl.debug_print(fmt, **kwargs) (any keyword arguments) inside a Pallas kernel compiled with jax.experimental.pallas with triton as the target; also passing a pytree of arguments that merge_callback_args flattens into kwargs.

Common situations: Porting a Pallas/Mosaic kernel from the TPU backend (where debug_print accepts kwargs) to the Triton GPU backend; using named arguments out of habit from Python's print-style APIs; upgrading JAX versions where the restriction started being enforced.

Related errors

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