{"record":{"id":"75e7c38a7bec8d19","repo":"jax-ml/jax","slug":"can-only-print-scalars-in-warp-level-code","errorCode":null,"errorMessage":"Can only print scalars in warp-level code.","messagePattern":"Can only print scalars in warp-level code\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3541,"sourceCode":"    raise ValueError(\n        \"Only positional arguments are supported by debug_print on Pallas.\"\n    )\n  primitives.check_debug_print_format(fmt, *args)\n  if not any(aval.shape for aval in ctx.avals_in):\n    scope = mgpu.ThreadSubset.WARPGROUP\n    if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n      scope = mgpu.ThreadSubset.WARP\n    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(","sourceCodeStart":3523,"sourceCodeEnd":3559,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3523-L3559","documentation":"Under warp-level semantics (PrimitiveSemantics.Warp), debug_print of a single non-scalar array is unsupported: warp-level code executes per-warp, so array printing is not implementable there.","triggerScenarios":"Calling jnp.debug_print/jax.debug.print with one array argument inside a kernel compiled with warp-level primitives (e.g. using warp-level ops from pallas GPU ops).","commonSituations":"Debugging kernels that mix warp-specialized code paths with prints; printing a block/tile value where the compiler picked Warp semantics.","solutions":["Print a scalar reduction of the array (e.g. jnp.sum(x)) instead of the array itself","Restructure so the print executes in warpgroup/lane semantics rather than warp semantics","Extract the tile to host memory and print outside the kernel"],"exampleFix":"// before\njax.debug.print(\"{}\", tile)  # in warp-level code\n// after\njax.debug.print(\"{}\", jnp.sum(tile))","handlingStrategy":"fallback","validationCode":"val = jnp.sum(x) if x.shape else x  # scalarize for warp-level print","typeGuard":"def printable_under_warp(aval) -> bool:\n    return aval.shape == ()","tryCatchPattern":null,"preventionTips":["Print scalars only in warp-level code","Reduce tiles to scalars before printing"],"tags":["jax","pallas","debug-print","warp-level"],"backgroundTag":"unsupported-feature-scope","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}