jax-ml/jax · error · ValueError

The kernel function in the pallas_call

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.

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

• BlockMapping for has captured constants
• Cannot lower a pallas_call with constants.
• Index map function for must not capture constants
• No constant handler for type
• a and b must be 2D, but got

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

Appendix: 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)