jax-ml/jax · error · ValueError

The emit_pipeline body function must return None.

Error message

The emit_pipeline body function must return None.

What it means

The pipelined kernel body is a loop body: it receives refs and (optionally) indices but must not return values. JAX traces the body and, if the flattened output tree has any leaves, raises this ValueError.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:2137

    # Trace with the global grid mapping to let the body resolve the mesh axes.
    with grid_mapping.trace_env():
      body_fun_dbg = api_util.debug_info(
          "emit_pipeline body", body, kernel_args, {}
      )
      in_avals_ft = ft.flatten(
          (kernel_args, {}),
          is_leaf=is_transformed_ref,
          registry=tracing_registry,
      )
      # Ensure the get_grid_mapping didn't produce TransformedRefs for tracing.
      assert all(
          not isinstance(x, state.TransformedRef) for x in in_avals_ft.vals)
      body_jaxpr, out_avals_ft = pe.trace_to_jaxpr(
          body, in_avals_ft, debug_info=body_fun_dbg
      )
      if out_avals_ft.tree.num_leaves != 0:
        raise ValueError("The emit_pipeline body function must return None.")

    all_index_map_consts = tuple(itertools.chain.from_iterable(
        bm.index_map_jaxpr.consts for bm in grid_mapping.block_mappings))

    refs_flat, refs_tree = tracing_registry.flatten(filtered_args)
    prim_args = EmitPipelinePrimitiveArgs(
        all_index_map_consts=all_index_map_consts,
        dynamic_grid_spec=dynamic_grid_specs,
        core_id=core_id,
        body_consts=tuple(body_jaxpr.consts),
        refs_flat=tuple(refs_flat),
        allocations=allocations,
    )
    args_flat, args_tree = tracing_registry.flatten(prim_args)
    return emit_pipeline_p.bind(
        *args_flat,
        grid_mapping=grid_mapping,
        body_jaxpr=body_jaxpr,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove all returns from the body; write outputs through the output refs in place
2. End the body with an explicit `return` (None) or a trailing no-op statement
3. Double-check lambdas: `lambda refs: refs.out.set(x)` returns None, `lambda refs: refs.out.set(x) or x` does not

Example fix

# before
def body(refs, i):
  refs.out[i] = compute(refs.inp[i])
  return refs.out  # WRONG
# after
def body(refs, i):
  refs.out[i] = compute(refs.inp[i])

assert body(refs_stub, idx_stub) is None, 'pipeline body must return None'

When it happens

Trigger: Writing the pipeline body with a bare or forgotten return, e.g. `return out_ref` or an expression body that evaluates to a value, instead of only performing in-place ref updates.

Common situations: Converting a pure-jax function (which returns outputs) into a Pallas pipeline body; copy-pasting a functional kernel where writes were returned rather than done via refs.

Related errors

• Sum of sizes {n} must be equal to dimension {axis} of the op
• Cannot pull iota along dimension {dimension} with None block
• Invalid memory space: {memory_space!r}
• dimension_semantics must be the same length as grid.
• Grid must consist of Python integers and JAX Arrays: {grid_t

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