{"record":{"id":"6475b7a976d68eaf","repo":"jax-ml/jax","slug":"ordered-debug-print-is-not-supported-on-pallas-6475b7","errorCode":null,"errorMessage":"Ordered debug_print is not supported on Pallas.","messagePattern":"Ordered debug_print is not supported on Pallas\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3520,"sourceCode":"@register_lowering_rule(\n    debugging.debug_print_p, mgpu.LoweringSemantics.Warpgroup\n)\n@register_lowering_rule(debugging.debug_print_p, *gpu_core.WGxWARP_SEMANTICS)\ndef _debug_print_lowering_rule(\n    ctx: LoweringRuleContext,\n    *args,\n    fmt,\n    ordered,\n    partitioned,\n    in_tree,\n    static_args,\n    np_printoptions,\n    has_placeholders,\n    logging_record,\n):\n  del partitioned, np_printoptions, has_placeholders\n  if ordered:\n    raise NotImplementedError(\"Ordered debug_print is not supported on Pallas.\")\n  args, kwargs = debugging.merge_callback_args(in_tree, args, static_args)\n  if kwargs:\n    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    )","sourceCodeStart":3502,"sourceCodeEnd":3538,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3502-L3538","documentation":"jax.debug.print(ordered=True) is not supported inside Pallas kernels on GPU. Ordered printing requires a deterministic callback ordering across devices/replicas which the Pallas lowering cannot provide.","triggerScenarios":"Calling jax.debug.print(..., ordered=True) inside a Pallas GPU kernel function.","commonSituations":"Copy-pasting debug prints from JIT-compiled JAX code (where ordered=True is common to keep output readable) into a Pallas kernel.","solutions":["Drop ordered=True and accept interleaved output","Use jax.debug.print without ordering, or print summary scalars after the kernel returns","For deterministic debugging, write values to an output buffer and inspect after launch"],"exampleFix":"// before\njax.debug.print(\"x={x}\", x=x, ordered=True)\n// after\njax.debug.print(\"x={x}\", x=x)","handlingStrategy":"validation","validationCode":"if ordered:\n    ordered = False  # not supported under Pallas","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid ordered=True in kernels","Print post-kernel scalars on host for ordered output"],"tags":["jax","pallas","debug-print","ordered"],"backgroundTag":"unsupported-feature-flag","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}