{"record":{"id":"407620678500ba36","repo":"jax-ml/jax","slug":"cannot-lower-a-pallas-call-with-constants","errorCode":null,"errorMessage":"Cannot lower a pallas_call with constants.","messagePattern":"Cannot lower a pallas_call with constants\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":848,"sourceCode":"    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.')\n  if interpret:\n    impl = partial(hlo_interpreter.pallas_call_hlo_interpret, **params)\n\n    try:\n      from jax._src.pallas.mosaic.interpret import interpret_pallas_call as mosaic_tpu_interpret  # pyrefly: ignore[missing-import]\n      from jax._src.pallas.mosaic.interpret import params as tpu_params  # pyrefly: ignore[missing-import]\n    except ImportError:\n      pass\n    else:\n      if isinstance(interpret, tpu_params.InterpretParams):\n        impl = partial(\n            mosaic_tpu_interpret.interpret_pallas_call,\n            interpret_params=interpret,\n            **params,\n        )\n\n    try:\n      from jax._src.pallas.mosaic_gpu.interpret import interpret_pallas_call as mosaic_gpu_interpret  # pyrefly: ignore[missing-import]","sourceCodeStart":830,"sourceCodeEnd":866,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L830-L866","documentation":"During MLIR lowering of pallas_call, JAX refuses to lower a jaxpr that carries constant variables. Pallas requires all arrays to flow in as runtime inputs (so BlockSpecs and the compiler can see them); consts in the jaxpr would be opaque to the Pallas compiler, so lowering raises ValueError.","triggerScenarios":"A pallas_call jaxpr with non-empty constvars, typically produced by a kernel that captured closed-over arrays (see the trace-time check) — e.g. constants reintroduced via intermediate transformations, caching of a traced kernel, or index/path transformations that re-add consts.","commonSituations":"Hitting the trace-time ValueError in error 2995 but bypassing it via an unusual path (custom primitives, jaxpr rewrite passes, pallas transforms); version changes where the trace-time check misses a case and the error surfaces only at lowering/compile time.","solutions":["Make every constant an explicit pallas_call input with its own BlockSpec","Audit the kernel for closures over jnp arrays and pass them as arguments","If the constant is small/static, bake it into the kernel as Python numbers so it does not become a constvar","Update JAX — newer versions catch this at trace time with a clearer message"],"exampleFix":"// before\npallas_call(kernel_closure_with_consts, out_shape=out)(x)  # jaxpr has constvars\n// after\ndef kernel(x_ref, w_ref, o_ref): ...\npallas_call(kernel, out_shape=out, in_specs=(spec_x, spec_w))(x, W)","handlingStrategy":"validation","validationCode":"jaxpr, consts = jax.make_jaxpr(kernel)(*ref_avals)\nif consts:\n    raise RuntimeError('pass constants as pallas_call inputs')","typeGuard":null,"tryCatchPattern":"try:\n    f = pallas_call(kernel, grid, out_shape); f.lower(x).compile()\nexcept ValueError as e:\n    if 'with constants' in str(e):\n        pass  # move constants to inputs and recompile","preventionTips":["Always compile-test kernels early (f.lower(...).compile()) to catch lowering errors","Keep kernels free of closed-over arrays"],"tags":["jax","pallas","lowering","constants","compiler"],"backgroundTag":"jit-captured-constant-tracing-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}