{"record":{"id":"13969948187dac61","repo":"jax-ml/jax","slug":"buffer-poisoning-is-not-supported-on-gpu-yet","errorCode":null,"errorMessage":"Buffer poisoning is not supported on GPU yet.","messagePattern":"Buffer poisoning is not supported on GPU yet\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3567,"sourceCode":"    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.\")\n  input_refs = []\n  should_discharge = []\n  wg_axis = ctx.module_ctx.axis_names.wg\n  is_multithreaded = wg_axis is not None\n  is_thread_collective = is_multithreaded and collective_axes == (wg_axis,)\n  # Make sure everyone has exited previous scoped allocations. Note that we\n  # don't synchronize when we exit the allocation, but only when we might want\n  # to reuse its memory again.\n  if collective_axes and collective_axes != (wg_axis,):\n    raise ValueError(\n        \"Only thread-collective allocations are supported in run_scoped.\"\n    )\n  if is_multithreaded and is_thread_collective:\n    gpu_dialect.barrier()\n  with contextlib.ExitStack() as alloc_stack:\n    for v in jaxpr.invars:\n      aval = cast(ShapedAbstractValue, v.aval)\n      if isinstance(aval, gpu_core.WGMMAAbstractAccumulatorRef):","sourceCodeStart":3549,"sourceCodeEnd":3585,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3549-L3585","documentation":"Buffer poisoning (a debug mode that fills output buffers with NaNs/garbage to catch uses of uninitialized memory, enabled via JAX flag pallas_default_buffer_poisoning / analogous config) is not implemented for the GPU Pallas backend, so run_scoped lowering refuses to proceed when it is on.","triggerScenarios":"Enabling buffer poisoning (e.g. jax.config.update('jax_default_buffer_poisoning_value', ...) or the pallas poisoning env/flag) and running a GPU kernel that uses pallas.run_scoped.","commonSituations":"Turning on aggressive memory debugging flags globally and then running the same suite on GPU; CI configs that enable poisoning on TPU-only.","solutions":["Disable buffer poisoning when running GPU kernels (set the poisoning value to None / off)","Gate the flag by platform: only enable for TUSA/TPU paths","Use other uninitialized-memory checks (e.g. compute-sanitizer) on GPU instead"],"exampleFix":"// before\njax.config.update('jax_default_buffer_poisoning_value', float('nan'))\n// after  # GPU\njax.config.update('jax_default_buffer_poisoning_value', None)","handlingStrategy":"validation","validationCode":"import jax\nif jax.config.jax_default_buffer_poisoning_value is not None and device == 'gpu':\n    jax.config.update('jax_default_buffer_poisoning_value', None)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Gate poisoning flags by platform","Use compute-sanitizer on GPU for uninit-memory checks"],"tags":["jax","pallas","buffer-poisoning","debug-flag","gpu"],"backgroundTag":"unsupported-debug-flag","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}