{"record":{"id":"5b0dce80e293262b","repo":"jax-ml/jax","slug":"debug-print-only-supports-printing-of-scalar-value","errorCode":null,"errorMessage":"debug_print only supports printing of scalar values, or a single array value when using the Mosaic GPU backend.","messagePattern":"debug_print only supports printing of scalar values, or a single array value when using the Mosaic GPU backend\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3549,"sourceCode":"    mgpu.debug_print(\n        fmt,\n        *(\n            _ensure_ir_value(arg, aval.dtype)\n            for arg, aval in zip(args, ctx.avals_in)\n        ),\n        scope=scope\n    )\n  elif len(ctx.avals_in) == 1:\n    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n      raise NotImplementedError(\"Can only print scalars in warp-level code.\")\n    [arg] = args\n    if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n      mgpu.dialect.debug_print(fmt, arg)\n    else:\n      arg.debug_print(fmt)\n\n  else:\n    raise NotImplementedError(\n        \"debug_print only supports printing of scalar values, or a single array\"\n        \" value when using the Mosaic GPU backend.\"\n    )\n\n  return ()\n\n\n@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(primitives.run_scoped_p, mgpu.LoweringSemantics.Warpgroup)\ndef _run_scoped_lowering_rule(\n    ctx: LoweringRuleContext,\n    *consts,\n    jaxpr: jax_core.Jaxpr,\n    collective_axes,\n    **_,\n):\n  if pallas_core.poison_buffers_enabled():\n    raise NotImplementedError(\"Buffer poisoning is not supported on GPU yet.\")","sourceCodeStart":3531,"sourceCodeEnd":3567,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3531-L3567","documentation":"Mosaic GPU's debug_print supports (a) any number of scalar arguments, or (b) exactly one array argument. Passing multiple arrays (or a mix where more than one input has shape) raises NotImplementedError.","triggerScenarios":"jax.debug.print(\"{} {}\", arr1, arr2) inside a Pallas GPU kernel where both args are arrays with non-empty shapes.","commonSituations":"Trying to print two tiles at once during kernel debugging; combining scalars and arrays in one call in the array branch.","solutions":["Split into separate debug_print calls, one per array","Or print scalars (e.g. per-tile sums) for each array in a single call","Remember the rule: many scalars OK, one array OK, multiple arrays not"],"exampleFix":"// before\njax.debug.print(\"{} {}\", a, b)\n// after\njax.debug.print(\"{}\", a)\njax.debug_print(\"{}\", b)","handlingStrategy":"validation","validationCode":"arrays = [a for a in args if a.shape]\nassert len(arrays) <= 1, 'one array max per debug_print'","typeGuard":"def debug_print_ok(avals) -> bool:\n    return sum(1 for a in avals if a.shape) <= 1","tryCatchPattern":null,"preventionTips":["One array per print call","Use scalar summaries when debugging multiple tiles"],"tags":["jax","pallas","debug-print","arrays"],"backgroundTag":"unsupported-feature-scope","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}