{"record":{"id":"8a55fc1d8009ca95","repo":"jax-ml/jax","slug":"2d-vectors-are-not-supported-in-debug-print-vec","errorCode":null,"errorMessage":"2D+ vectors are not supported in debug_print: {vec_ty}","messagePattern":"2D\\+ vectors are not supported in debug_print: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":195,"sourceCode":"    )\n  if isinstance(arg.type, (ir.BF16Type, ir.F16Type)):\n    arg = arith.extf(ir.F32Type.get(), arg)\n    return \"%f\", arg\n  raise NotImplementedError(f\"Can't print the type {arg.type}\")\n\n\ndef debug_print(fmt, *args, uniform=True, scope=None):\n  if not uniform and scope is not None:\n    raise ValueError(\"Cannot specify scope to a non-uniform debug_print.\")\n  if scope is None:\n    scope = ThreadSubset.WARPGROUP\n  type_formats = []\n  new_args = []\n  for arg in args:\n    if isinstance(arg.type, ir.VectorType):\n      vec_ty = ir.VectorType(arg.type)\n      if len(vec_ty.shape) > 1:\n        raise NotImplementedError(\n            f\"2D+ vectors are not supported in debug_print: {vec_ty}\"\n        )\n      vec_args = [\n          vector.extract(\n              arg,\n              dynamic_position=[],\n              static_position=ir.DenseI64ArrayAttr.get([i]),\n          )\n          for i in range(vec_ty.shape[0])\n      ]\n      ty_formats, args = zip(*map(_debug_scalar_ty_format, vec_args))\n      ty_format = f\"[{','.join(ty_formats)}]\"\n      new_args += args\n    else:\n      ty_format, arg = _debug_scalar_ty_format(arg)\n      new_args.append(arg)\n\n    if ty_format is None:","sourceCodeStart":177,"sourceCodeEnd":213,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L177-L213","documentation":"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.","triggerScenarios":"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.","commonSituations":"Debugging tcgen05/wgmma kernels that keep 2D register tiles; after refactors that changed scalars to tiled vectors passed into debug prints.","solutions":["Flatten before printing: arith.reshape / vector.shape_cast the 2D vector to 1D, then pass the 1D vector","Or extract a specific row with vector.extract (dynamic_position=[row]) and print that 1D vector","Print individual scalars via vector.extract with full positions in a loop over a few representative lanes"],"exampleFix":"# before\ndebug_print('acc={}', acc_2d)  # vector<8x8xf32>\n# after\nflat = vector.shape_cast(acc_2d, ir.VectorType.get((64,), f32))\ndebug_print('acc={}', flat)","handlingStrategy":"fallback","validationCode":"if isinstance(arg.type, ir.VectorType) and len(arg.type.shape) > 1:\n    arg = vector.shape_cast(arg, ir.VectorType.get((int(np.prod(arg.type.shape)),), arg.type.element_type))\ndebug_print('m={}', arg)","typeGuard":"def is_printable_vector(arg):\n    return not isinstance(arg.type, ir.VectorType) or len(arg.type.shape) <= 1","tryCatchPattern":null,"preventionTips":["Flatten 2D register tiles with vector.shape_cast before debug_print","Extract single rows for targeted inspection instead of whole fragments"],"tags":["gpu","mosaic","debug-print","vector","rank"],"backgroundTag":"unsupported-vector-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}