{"record":{"id":"4a82eb33c7b54bfe","repo":"jax-ml/jax","slug":"the-emit-pipeline-body-function-must-return-none","errorCode":null,"errorMessage":"The emit_pipeline body function must return None.","messagePattern":"The emit_pipeline body function must return None\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":2137,"sourceCode":"\n    # Trace with the global grid mapping to let the body resolve the mesh axes.\n    with grid_mapping.trace_env():\n      body_fun_dbg = api_util.debug_info(\n          \"emit_pipeline body\", body, kernel_args, {}\n      )\n      in_avals_ft = ft.flatten(\n          (kernel_args, {}),\n          is_leaf=is_transformed_ref,\n          registry=tracing_registry,\n      )\n      # Ensure the get_grid_mapping didn't produce TransformedRefs for tracing.\n      assert all(\n          not isinstance(x, state.TransformedRef) for x in in_avals_ft.vals)\n      body_jaxpr, out_avals_ft = pe.trace_to_jaxpr(\n          body, in_avals_ft, debug_info=body_fun_dbg\n      )\n      if out_avals_ft.tree.num_leaves != 0:\n        raise ValueError(\"The emit_pipeline body function must return None.\")\n\n    all_index_map_consts = tuple(itertools.chain.from_iterable(\n        bm.index_map_jaxpr.consts for bm in grid_mapping.block_mappings))\n\n    refs_flat, refs_tree = tracing_registry.flatten(filtered_args)\n    prim_args = EmitPipelinePrimitiveArgs(\n        all_index_map_consts=all_index_map_consts,\n        dynamic_grid_spec=dynamic_grid_specs,\n        core_id=core_id,\n        body_consts=tuple(body_jaxpr.consts),\n        refs_flat=tuple(refs_flat),\n        allocations=allocations,\n    )\n    args_flat, args_tree = tracing_registry.flatten(prim_args)\n    return emit_pipeline_p.bind(\n        *args_flat,\n        grid_mapping=grid_mapping,\n        body_jaxpr=body_jaxpr,","sourceCodeStart":2119,"sourceCodeEnd":2155,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L2119-L2155","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove all returns from the body; write outputs through the output refs in place","End the body with an explicit `return` (None) or a trailing no-op statement","Double-check lambdas: `lambda refs: refs.out.set(x)` returns None, `lambda refs: refs.out.set(x) or x` does not"],"exampleFix":"# before\ndef body(refs, i):\n  refs.out[i] = compute(refs.inp[i])\n  return refs.out  # WRONG\n# after\ndef body(refs, i):\n  refs.out[i] = compute(refs.inp[i])","handlingStrategy":"validation","validationCode":"assert body(refs_stub, idx_stub) is None, 'pipeline body must return None'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Bodies must write through refs only; end with `return`","Watch out for lambdas whose last expression returns a value"],"tags":["jax","pallas","kernel-body","validation"],"backgroundTag":"unexpected-return-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}