jax-ml/jax · error · NotImplementedError

Can't print the type {arg.type}

Error message

Can't print the type {arg.type}

What it means

Raised by _debug_scalar_ty_format, the helper behind debug_print, when asked to produce a printf format string for a scalar type it doesn't know how to print. Currently supported: index/integer (via %llu with i32/i8 handling), f32/f64, and bf16/f16 (extended to f32 and printed with %f). Any other scalar (e.g. f8, complex, bool-like custom types) has no format mapping.

Source

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

def _debug_scalar_ty_format(arg):
  if isinstance(arg.type, ir.IndexType):
    return "%llu", arg
  if isinstance(arg.type, ir.IntegerType):
    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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Bitcast the value to a printable width first (e.g. bitcast f8 -> i8, or extend to f32 where semantics allow) and print as integer/float
  2. Print a numeric proxy: convert to f32 via arith.extf only for types that support it, else bitcast to i32
  3. Remove the debug_print for unsupported dtypes and inspect via IR dumps or dumps to memory instead

Example fix

# before
debug_print('x={}', f8_val)  # NotImplementedError
# after
x_bits = arith.bitcast(i8_ty, f8_val)
debug_print('x_bits={}', x_bits)
Defensive patterns

Strategy: fallback

Validate before calling

supported = (ir.IndexType, ir.IntegerType, ir.F32Type, ir.F64Type, ir.BF16Type, ir.F16Type)
if not isinstance(arg.type, supported):
    arg = arith.bitcast(ir.IntegerType.get_signless(32), arg)  # print raw bits
debug_print('v={}', arg)

Type guard

def is_printable_scalar(arg):
    return isinstance(arg.type, (ir.IndexType, ir.IntegerType, ir.F32Type, ir.F64Type, ir.BF16Type, ir.F16Type))

Prevention

When it happens

Trigger: Calling utils.debug_print('fmt {}', arg) where arg is an f8e4m3fn value, a complex value, or any scalar type outside index/integer/float16-family/f32/f64.

Common situations: Debugging FP8 matmul kernels on Blackwell and trying to print raw operand values; passing values that were bitcast to unusual integer widths (e.g. i4-packed) before printing.

Related errors


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