jax-ml/jax · error · ValueError

The kernel function in the pallas_call {debug_info.func_src_

Error message

The kernel function in the pallas_call {debug_info.func_src_info} captures constants [{pp_consts_avals}]. You should pass them as inputs.

What it means

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.

Source

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

    closed_jaxpr, out_avals = pe.trace_to_jaxpr(
        fun_with_transforms, kernel_avals,
        debug_info)
    consts = closed_jaxpr.consts
    jaxpr, _ = pe.dce_jaxpr(closed_jaxpr,
                            used_outputs=[True] * len(closed_jaxpr.outvars),
                            instantiate=True)
    if consts:
      consts_avals = [
          aval
          for c in consts
          if not isinstance(aval := jax_core.typeof(c), state.AbstractRef)
      ]
      if consts_avals:
        ctx = jax_core.JaxprPpContext()
        pp_consts_avals = ", ".join(
            jax_core.pp_aval(aval, ctx) for aval in consts_avals
        )
        raise ValueError(
            "The kernel function in the pallas_call"
            f" {debug_info.func_src_info} captures constants"
            f" [{pp_consts_avals}]. You should pass them as inputs."
        )

  kernel_out_tree = out_avals.tree
  if not indexer and kernel_out_tree != tree_util.tree_structure(None):
    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."

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add the captured array as an explicit parameter of the kernel function and pass it through pallas_call(...)(x, W) with a matching BlockSpec
2. 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
3. Re-run and confirm the error's listed avals ([...]) are now absent

Example fix

// before
def kernel(x_ref, o_ref):
  o_ref[...] = x_ref[...] + W  # W captured from outer scope
pallas_call(kernel, out_shape=out)(x)
// after
def kernel(x_ref, w_ref, o_ref):
  o_ref[...] = x_ref[...] + w_ref[...]
pallas_call(kernel, out_shape=out, in_specs=[BlockSpec(...), BlockSpec(...) ])(x, W)

# Trace the kernel standalone to detect captures before pallas_call
jaxpr = jax.make_jaxpr(kernel)(refs...)
assert not jaxpr.jaxpr.constvars, f'kernel captures: {jaxpr.jaxpr.constvars}'

def kernel_is_closure_free(kernel, *refs) -> bool:
    jaxpr = jax.make_jaxpr(kernel)(*refs)
    return not jaxpr.jaxpr.constvars

try:
    pallas_call(kernel, grid, out_shape)(x)
except ValueError as e:
    if 'captures constants' in str(e):
        # move the listed avals to explicit inputs

When it happens

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

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

Related errors

• No constant handler for type: {type(val)}
• Index map function {debug_info.func_src_info} for {origin} m
• BlockMapping for {self.origin} has captured constants: {self
• Cannot lower a pallas_call with constants.
• The unsafe_buffer_pointer() method was called on {self._erro

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