{"record":{"id":"acaf43219dd32393","repo":"jax-ml/jax","slug":"non-debug-checks-are-not-supported-by-the-mosaic-g","errorCode":null,"errorMessage":"Non-debug checks are not supported by the Mosaic GPU backend. Functionalize them via `jax.experimental.checkify`.","messagePattern":"Non-debug checks are not supported by the Mosaic GPU backend\\. Functionalize them via `jax\\.experimental\\.checkify`\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":4606,"sourceCode":"    )\n    block = custom_op.body.blocks.append(*[o.type for o in operands])\n    with ir.InsertionPoint(block):\n      ctx.module_ctx.lowering_semantics = mgpu.LoweringSemantics.Lane\n      try:\n        yield list(block.arguments)\n      finally:\n        ctx.module_ctx.lowering_semantics = mgpu.LoweringSemantics.Warpgroup\n      mgpu.dialect.ReturnOp(operands_=[])\n    _isolate_from_above(custom_op)\n  else:\n    yield list(operands)\n\n\n@register_lowering_rule(checkify.check_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(checkify.check_p, mgpu.LoweringSemantics.Warpgroup)\ndef _check_lowering_rule(ctx: LoweringRuleContext, *err_args, err_tree, debug):\n  if not debug:\n    raise NotImplementedError(\n        \"Non-debug checks are not supported by the Mosaic GPU backend.\"\n        \" Functionalize them via `jax.experimental.checkify`.\"\n    )\n  if not pallas_core.debug_checks_enabled():\n    return []\n\n  error = jax.tree.unflatten(err_tree, err_args)\n  [pred] = error._pred.values()\n  [exception_tree] = error._metadata.values()\n  [payload] = error._payload.values()\n  exception = jax.tree.unflatten(exception_tree, payload)\n  assert isinstance(exception, checkify.FailedCheckError)\n\n  # check_p has an inverted predicate compared to assert, so we need to compute\n  # ``not pred`` here.\n  minus_one = _ir_constant(-1, mgpu_utils.dtype_to_ir_type(jnp.bool))\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    pred = pred.registers.item()","sourceCodeStart":4588,"sourceCodeEnd":4624,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L4588-L4624","documentation":"checkify.check_p lowered with debug=False is not supported on the Mosaic GPU backend: functionalized checks cannot run on GPU without checkify's transformation, so only debug-mode checks are allowed.","triggerScenarios":"Using checkify.check (non-debug) inside a pallas kernel compiled for Mosaic GPU without applying checkify.functionalize to the whole computation.","commonSituations":"Adding runtime assertions to GPU kernels and forgetting to wrap the outer function with checkify.checkify / functionalize; passing debug=False to check.","solutions":["Wrap the kernel launch with checkify.functionalize and run the error payload check afterwards","Or use debug checks (they no-op unless JAX_DEBUG_CHECKS / debug_checks is enabled)","Avoid check inside pallas kernels on GPU if functionalization is impractical"],"exampleFix":"# before\n@pl.kernel\n def k(ref): checkify.check(cond, \"msg\")\njax_result = plgpu_kernel(...)\n# after\nchecked = checkify.functionalize(plgpu_kernel)\nerr, jax_result = checked(...)\ncheckify.check_error(err)","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":["Wrap kernel launches with checkify.functionalize and check_error","Or run with debug checks enabled (JAX_DEBUG_CHECKS)"],"tags":["pallas","checkify","mosaic-gpu","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}