{"record":{"id":"b6157896f1572b5a","repo":"jax-ml/jax","slug":"the-kernel-function-in-the-pallas-call-debug-info-b61578","errorCode":null,"errorMessage":"The kernel function in the pallas_call {debug_info.func_src_info} should return None. It returns a PyTree: {kernel_out_tree}","messagePattern":"The kernel function in the pallas_call (.+?) should return None\\. It returns a PyTree: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":830,"sourceCode":"      consts_avals = [\n          aval\n          for c in consts\n          if not isinstance(aval := jax_core.typeof(c), state.AbstractRef)\n      ]\n      if consts_avals:\n        ctx = jax_core.JaxprPpContext()\n        pp_consts_avals = \", \".join(\n            jax_core.pp_aval(aval, ctx) for aval in consts_avals\n        )\n        raise ValueError(\n            \"The kernel function in the pallas_call\"\n            f\" {debug_info.func_src_info} captures constants\"\n            f\" [{pp_consts_avals}]. You should pass them as inputs.\"\n        )\n\n  kernel_out_tree = out_avals.tree\n  if not indexer and kernel_out_tree != tree_util.tree_structure(None):\n    raise ValueError(\n        f\"The kernel function in the pallas_call {debug_info.func_src_info} \"\n        f\"should return None. It returns a PyTree: {kernel_out_tree}\")\n  return jaxpr, tuple(consts)\n\n\ndef _unsupported_lowering_error(platform: str) -> Exception:\n  return ValueError(\n      f\"Cannot lower pallas_call on platform: {platform}. To use Pallas on GPU,\"\n      \" install jaxlib GPU. To use Pallas on TPU, install jaxlib TPU and\"\n      \" libtpu. See https://docs.jax.dev/en/latest/installation.html.\"\n  )\n\n\ndef _pallas_call_lowering(\n    ctx: mlir.LoweringRuleContext, *in_nodes, interpret: Any, **params\n):\n  if params['jaxpr'].constvars:\n    raise ValueError('Cannot lower a pallas_call with constants.')","sourceCodeStart":812,"sourceCodeEnd":848,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L812-L848","documentation":"Pallas kernels communicate results exclusively through output reference arguments (o_ref); the kernel function itself must return None. After tracing, if the kernel's return pytree structure is anything other than None (non-indexer mode), JAX raises this ValueError showing the actual returned tree structure.","triggerScenarios":"Writing a Pallas kernel with a return statement, e.g. `return x_ref[...] + 1`, or returning a tuple of results instead of assigning them into the output refs.","commonSituations":"Developers coming from jit/vmap style where functions return values; converting a pure-jnp function into a Pallas kernel and keeping the return statement.","solutions":["Delete the return statement and write results into the output reference(s): o_ref[...] = result","Verify out_shape matches the number of output refs declared in the kernel signature","If returning multiple results, declare multiple o_ref parameters and assign each one"],"exampleFix":"// before\ndef kernel(x_ref, o_ref):\n  return jnp.exp(x_ref[...])  # raises\n// after\ndef kernel(x_ref, o_ref):\n  o_ref[...] = jnp.exp(x_ref[...])","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def kernel_returns_none(kernel) -> bool:\n    return inspect.isfunction(kernel) and (\n        # static check: kernel writes to *_ref params; runtime check via interpret mode\n        True if _run_interpret_returns_none(kernel) else False)","tryCatchPattern":"try:\n    pallas_call(kernel, grid, out_shape)(x)\nexcept ValueError as e:\n    if 'should return None' in str(e):\n        kernel = strip_return(kernel)  # rewrite to assign into o_ref","preventionTips":["Pallas kernels always write results into o_ref; never return","Test kernels with interpret=True first for fast validation"],"tags":["jax","pallas","kernel","return-value","validation"],"backgroundTag":"kernel-must-write-outputs-not-return","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}