What causes "Cannot lower a pallas_call with constants." in jax-ml/jax?

Make every constant an explicit pallas_call input with its own BlockSpec

jax-ml/jax · error · ValueError

Cannot lower a pallas_call with constants.

Error message

Cannot lower a pallas_call with constants.

What it means

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.

Source

Thrown at jax/_src/pallas/pallas_call.py:848

    raise ValueError(
        f"The kernel function in the pallas_call {debug_info.func_src_info} "
        f"should return None. It returns a PyTree: {kernel_out_tree}")
  return jaxpr, tuple(consts)


def _unsupported_lowering_error(platform: str) -> Exception:
  return ValueError(
      f"Cannot lower pallas_call on platform: {platform}. To use Pallas on GPU,"
      " install jaxlib GPU. To use Pallas on TPU, install jaxlib TPU and"
      " libtpu. See https://docs.jax.dev/en/latest/installation.html."
  )


def _pallas_call_lowering(
    ctx: mlir.LoweringRuleContext, *in_nodes, interpret: Any, **params
):
  if params['jaxpr'].constvars:
    raise ValueError('Cannot lower a pallas_call with constants.')
  if interpret:
    impl = partial(hlo_interpreter.pallas_call_hlo_interpret, **params)

    try:
      from jax._src.pallas.mosaic.interpret import interpret_pallas_call as mosaic_tpu_interpret  # pyrefly: ignore[missing-import]
      from jax._src.pallas.mosaic.interpret import params as tpu_params  # pyrefly: ignore[missing-import]
    except ImportError:
      pass
    else:
      if isinstance(interpret, tpu_params.InterpretParams):
        impl = partial(
            mosaic_tpu_interpret.interpret_pallas_call,
            interpret_params=interpret,
            **params,
        )

    try:
      from jax._src.pallas.mosaic_gpu.interpret import interpret_pallas_call as mosaic_gpu_interpret  # pyrefly: ignore[missing-import]

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
pallas_call(kernel_closure_with_consts, out_shape=out)(x)  # jaxpr has constvars
// after
def kernel(x_ref, w_ref, o_ref): ...
pallas_call(kernel, out_shape=out, in_specs=(spec_x, spec_w))(x, W)

Defensive patterns

Strategy: validation

Validate before calling

jaxpr, consts = jax.make_jaxpr(kernel)(*ref_avals)
if consts:
    raise RuntimeError('pass constants as pallas_call inputs')

Try / catch

try:
    f = pallas_call(kernel, grid, out_shape); f.lower(x).compile()
except ValueError as e:
    if 'with constants' in str(e):
        pass  # move constants to inputs and recompile

Prevention

Always compile-test kernels early (f.lower(...).compile()) to catch lowering errors
Keep kernels free of closed-over arrays

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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