jax-ml/jax · error · NotImplementedError

2D+ vectors are not supported in debug_print: {vec_ty}

Error message

2D+ vectors are not supported in debug_print: {vec_ty}

What it means

Raised by debug_print when one of the arguments is a vector of rank 2 or higher. The helper unrolls vector arguments by extracting each lane with vector.extract, which only supports 1D positions, so 2D+ vectors (e.g. a 8x8 f32 fragment) cannot be printed element-wise through this path.

Source

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

    )
  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])
      ]
      ty_formats, args = zip(*map(_debug_scalar_ty_format, vec_args))
      ty_format = f"[{','.join(ty_formats)}]"
      new_args += args
    else:
      ty_format, arg = _debug_scalar_ty_format(arg)
      new_args.append(arg)

    if ty_format is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten before printing: arith.reshape / vector.shape_cast the 2D vector to 1D, then pass the 1D vector
2. Or extract a specific row with vector.extract (dynamic_position=[row]) and print that 1D vector
3. Print individual scalars via vector.extract with full positions in a loop over a few representative lanes

Example fix

# before
debug_print('acc={}', acc_2d)  # vector<8x8xf32>
# after
flat = vector.shape_cast(acc_2d, ir.VectorType.get((64,), f32))
debug_print('acc={}', flat)

if isinstance(arg.type, ir.VectorType) and len(arg.type.shape) > 1:
    arg = vector.shape_cast(arg, ir.VectorType.get((int(np.prod(arg.type.shape)),), arg.type.element_type))
debug_print('m={}', arg)

def is_printable_vector(arg):
    return not isinstance(arg.type, ir.VectorType) or len(arg.type.shape) <= 1

When it happens

Trigger: Calling utils.debug_print('m={}', v) where v has ir.VectorType like vector<8x8xf32> — typical when inspecting accumulator fragments or register tiles in an MMA pipeline.

Common situations: Debugging tcgen05/wgmma kernels that keep 2D register tiles; after refactors that changed scalars to tiled vectors passed into debug prints.

Related errors

• `strides` must contain only 1s.
• Only slicing with static indices allowed.
• Unsupported reduction kind: {op.kind}
• Can't print the type {arg.type}
• Cannot specify scope to a non-uniform debug_print.

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