jax-ml/jax · error · NotImplementedError

Non-debug checks are not supported by the Mosaic GPU backend

Error message

Non-debug checks are not supported by the Mosaic GPU backend. Functionalize them via `jax.experimental.checkify`.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:4606

    )
    block = custom_op.body.blocks.append(*[o.type for o in operands])
    with ir.InsertionPoint(block):
      ctx.module_ctx.lowering_semantics = mgpu.LoweringSemantics.Lane
      try:
        yield list(block.arguments)
      finally:
        ctx.module_ctx.lowering_semantics = mgpu.LoweringSemantics.Warpgroup
      mgpu.dialect.ReturnOp(operands_=[])
    _isolate_from_above(custom_op)
  else:
    yield list(operands)


@register_lowering_rule(checkify.check_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(checkify.check_p, mgpu.LoweringSemantics.Warpgroup)
def _check_lowering_rule(ctx: LoweringRuleContext, *err_args, err_tree, debug):
  if not debug:
    raise NotImplementedError(
        "Non-debug checks are not supported by the Mosaic GPU backend."
        " Functionalize them via `jax.experimental.checkify`."
    )
  if not pallas_core.debug_checks_enabled():
    return []

  error = jax.tree.unflatten(err_tree, err_args)
  [pred] = error._pred.values()
  [exception_tree] = error._metadata.values()
  [payload] = error._payload.values()
  exception = jax.tree.unflatten(exception_tree, payload)
  assert isinstance(exception, checkify.FailedCheckError)

  # check_p has an inverted predicate compared to assert, so we need to compute
  # ``not pred`` here.
  minus_one = _ir_constant(-1, mgpu_utils.dtype_to_ir_type(jnp.bool))
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    pred = pred.registers.item()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the kernel launch with checkify.functionalize and run the error payload check afterwards
2. Or use debug checks (they no-op unless JAX_DEBUG_CHECKS / debug_checks is enabled)
3. Avoid check inside pallas kernels on GPU if functionalization is impractical

Example fix

# before
@pl.kernel
 def k(ref): checkify.check(cond, "msg")
jax_result = plgpu_kernel(...)
# after
checked = checkify.functionalize(plgpu_kernel)
err, jax_result = checked(...)
checkify.check_error(err)

When it happens

Trigger: Using checkify.check (non-debug) inside a pallas kernel compiled for Mosaic GPU without applying checkify.functionalize to the whole computation.

Common situations: Adding runtime assertions to GPU kernels and forgetting to wrap the outer function with checkify.checkify / functionalize; passing debug=False to check.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Sparse metadata format not implemented for {operand_dtype=}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/acaf43219dd32393. Report an issue: GitHub.