{"record":{"id":"3ad526e218ccfa78","repo":"jax-ml/jax","slug":"the-kernel-function-in-the-pallas-call-debug-info","errorCode":null,"errorMessage":"The kernel function in the pallas_call {debug_info.func_src_info} captures constants [{pp_consts_avals}]. You should pass them as inputs.","messagePattern":"The kernel function in the pallas_call (.+?) captures constants \\[(.+?)\\]\\. You should pass them as inputs\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":822,"sourceCode":"    closed_jaxpr, out_avals = pe.trace_to_jaxpr(\n        fun_with_transforms, kernel_avals,\n        debug_info)\n    consts = closed_jaxpr.consts\n    jaxpr, _ = pe.dce_jaxpr(closed_jaxpr,\n                            used_outputs=[True] * len(closed_jaxpr.outvars),\n                            instantiate=True)\n    if consts:\n      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.\"","sourceCodeStart":804,"sourceCodeEnd":840,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L804-L840","documentation":"When JAX traces the Pascal kernel function to a jaxpr, any array the kernel closes over from the enclosing Python scope becomes a constant of the traced jaxpr. pallas_call requires every array used by a kernel to be an explicit input (it must appear in the call signature and BlockSpecs), so captured constants trigger this ValueError listing the offending avals.","triggerScenarios":"Defining a kernel like def kernel(x_ref): compute(x_ref, W) where W is an ndarray from the outer scope; passing a jnp array into the kernel via closure instead of via pallas_call arguments.","commonSituations":"Writing Pallas kernels that reference module-level weights or hyperparameter arrays; refactoring a jit function into a Pallas kernel and forgetting to thread constants through as inputs.","solutions":["Add the captured array as an explicit parameter of the kernel function and pass it through pallas_call(...)(x, W) with a matching BlockSpec","If the value is a true compile-time constant, make it a Python scalar/static value baked into the kernel body rather than a traced jnp array","Re-run and confirm the error's listed avals ([...]) are now absent"],"exampleFix":"// before\ndef kernel(x_ref, o_ref):\n  o_ref[...] = x_ref[...] + W  # W captured from outer scope\npallas_call(kernel, out_shape=out)(x)\n// after\ndef kernel(x_ref, w_ref, o_ref):\n  o_ref[...] = x_ref[...] + w_ref[...]\npallas_call(kernel, out_shape=out, in_specs=[BlockSpec(...), BlockSpec(...) ])(x, W)","handlingStrategy":"type-guard","validationCode":"# Trace the kernel standalone to detect captures before pallas_call\njaxpr = jax.make_jaxpr(kernel)(refs...)\nassert not jaxpr.jaxpr.constvars, f'kernel captures: {jaxpr.jaxpr.constvars}'","typeGuard":"def kernel_is_closure_free(kernel, *refs) -> bool:\n    jaxpr = jax.make_jaxpr(kernel)(*refs)\n    return not jaxpr.jaxpr.constvars","tryCatchPattern":"try:\n    pallas_call(kernel, grid, out_shape)(x)\nexcept ValueError as e:\n    if 'captures constants' in str(e):\n        # move the listed avals to explicit inputs","preventionTips":["Never reference outer-scope jnp arrays inside Pallas kernels","Pass weights/constants as kernel parameters with BlockSpecs"],"tags":["jax","pallas","closure","constants","tracing"],"backgroundTag":"jit-captured-constant-tracing-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}