jax-ml/jax · error · NotImplementedError

Can only print scalars in warp-level code.

Error message

Can only print scalars in warp-level code.

What it means

Under warp-level semantics (PrimitiveSemantics.Warp), debug_print of a single non-scalar array is unsupported: warp-level code executes per-warp, so array printing is not implementable there.

Source

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

    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.")
    [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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print a scalar reduction of the array (e.g. jnp.sum(x)) instead of the array itself
2. Restructure so the print executes in warpgroup/lane semantics rather than warp semantics
3. Extract the tile to host memory and print outside the kernel

Example fix

// before
jax.debug.print("{}", tile)  # in warp-level code
// after
jax.debug.print("{}", jnp.sum(tile))

val = jnp.sum(x) if x.shape else x  # scalarize for warp-level print

def printable_under_warp(aval) -> bool:
    return aval.shape == ()

When it happens

Trigger: Calling jnp.debug_print/jax.debug.print with one array argument inside a kernel compiled with warp-level primitives (e.g. using warp-level ops from pallas GPU ops).

Common situations: Debugging kernels that mix warp-specialized code paths with prints; printing a block/tile value where the compiler picked Warp semantics.

Related errors

• Ordered debug_print is not supported on Pallas.
• Only positional arguments are supported by debug_print on Pa
• 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/75e7c38a7bec8d19. Report an issue: GitHub.