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} should return None. It returns a PyTree: {kernel_out_tree}

What it means

Pallas kernels communicate results exclusively through output reference arguments (o_ref); the kernel function itself must return None. After tracing, if the kernel's return pytree structure is anything other than None (non-indexer mode), JAX raises this ValueError showing the actual returned tree structure.

Source

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

      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."
  )


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.')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Delete the return statement and write results into the output reference(s): o_ref[...] = result
Verify out_shape matches the number of output refs declared in the kernel signature
If returning multiple results, declare multiple o_ref parameters and assign each one

Example fix

// before
def kernel(x_ref, o_ref):
  return jnp.exp(x_ref[...])  # raises
// after
def kernel(x_ref, o_ref):
  o_ref[...] = jnp.exp(x_ref[...])

Defensive patterns

Strategy: type-guard

Type guard

def kernel_returns_none(kernel) -> bool:
    return inspect.isfunction(kernel) and (
        # static check: kernel writes to *_ref params; runtime check via interpret mode
        True if _run_interpret_returns_none(kernel) else False)

Try / catch

try:
    pallas_call(kernel, grid, out_shape)(x)
except ValueError as e:
    if 'should return None' in str(e):
        kernel = strip_return(kernel)  # rewrite to assign into o_ref

Prevention

Pallas kernels always write results into o_ref; never return
Test kernels with interpret=True first for fast validation

When it happens

Trigger: Writing a Pallas kernel with a return statement, e.g. `return x_ref[...] + 1`, or returning a tuple of results instead of assigning them into the output refs.

Common situations: Developers coming from jit/vmap style where functions return values; converting a pure-jnp function into a Pallas kernel and keeping the return statement.

Related errors

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