jax-ml/jax · error · ValueError

Cannot specify scope to a non-uniform debug_print.

Error message

Cannot specify scope to a non-uniform debug_print.

What it means

Raised by debug_print when the caller passes scope=... together with uniform=False. A 'non-uniform' debug print executes in every thread independently, so a single-threaded scope restriction only makes sense for uniform prints; mixing them is contradictory and rejected.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:186

    if ir.IntegerType(arg.type).width < 64:
      arg = arith.extui(ir.IntegerType.get_signless(64), arg)
    return "%llu", arg
  if isinstance(arg.type, ir.F32Type):
    return "%f", arg
  if isinstance(arg.type, ir.Float8E8M0FNUType):
    return "%u", arith.extui(
        ir.IntegerType.get_signless(32),
        arith.bitcast(ir.IntegerType.get_signless(8), arg),
    )
  if isinstance(arg.type, (ir.BF16Type, ir.F16Type)):
    arg = arith.extf(ir.F32Type.get(), arg)
    return "%f", arg
  raise NotImplementedError(f"Can't print the type {arg.type}")


def debug_print(fmt, *args, uniform=True, scope=None):
  if not uniform and scope is not None:
    raise ValueError("Cannot specify scope to a non-uniform debug_print.")
  if scope is None:
    scope = ThreadSubset.WARPGROUP
  type_formats = []
  new_args = []
  for arg in args:
    if isinstance(arg.type, ir.VectorType):
      vec_ty = ir.VectorType(arg.type)
      if len(vec_ty.shape) > 1:
        raise NotImplementedError(
            f"2D+ vectors are not supported in debug_print: {vec_ty}"
        )
      vec_args = [
          vector.extract(
              arg,
              dynamic_position=[],
              static_position=ir.DenseI64ArrayAttr.get([i]),
          )
          for i in range(vec_ty.shape[0])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Drop the scope argument when using uniform=False
Or keep scope and use uniform=True (default) if you only want one thread per scope to print

Example fix

# before
debug_print('v={}', v, uniform=False, scope=ThreadSubset.WARPGROUP)
# after
debug_print('v={}', v, uniform=False)

Defensive patterns

Strategy: validation

Validate before calling

if not uniform and scope is not None:
    scope = None  # scope only applies to uniform prints
debug_print(fmt, *args, uniform=uniform, scope=scope)

Prevention

Pass scope only with uniform=True (the default)
Build small wrapper macros that encode the uniform/scope policy once

When it happens

Trigger: Calling utils.debug_print(fmt, *args, uniform=False, scope=ThreadSubset.WARP) — any combination where uniform is False and scope is not None.

Common situations: Copy-pasting a uniform debug_print call and flipping uniform=False to see per-thread output while leaving the scope kwarg in place; refactoring shared debug macros that always pass scope.

Related errors

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