{"record":{"id":"357c35f417ac7e27","repo":"jax-ml/jax","slug":"the-kernel-function-in-mpmd-map-debug-info-func-s-357c35","errorCode":null,"errorMessage":"The kernel function in mpmd_map {debug_info.func_src_info} should return None. It returns a PyTree: {fun_out_tree}.","messagePattern":"The kernel function in mpmd_map (.+?) should return None\\. It returns a PyTree: (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mpmd.py","lineNumber":986,"sourceCode":"    consts_per_fn = []\n    debug_infos = [api_util.debug_info(\"mpmd_map\", fn, kernel_arg_avals, kernel_kwarg_avals)\n                   for _, fn in meshes_and_fns]\n    if name is not None:\n      debug_infos = [di.replace_func_name(name) for di in debug_infos]\n    # If names are non-distinct (e.g. because user passed multiple functions\n    # with the same name, or because of the name= arg handled just above),\n    # uniquify them with the core type.\n    if len({di.func_name for di in debug_infos}) != len(debug_infos):\n      debug_infos = [di.replace_func_name(f\"{di.func_name}__{mesh.core_type}\")\n                     for di, mesh in zip(debug_infos, meshes)]\n    for (mesh, fn), debug_info in zip(meshes_and_fns, debug_infos):\n      with mpmd_map_tracing_context(mesh, all_meshes):\n        jaxpr, out_avals = pe.trace_to_jaxpr(\n            fn, in_avals_ft, debug_info\n        )\n      fun_out_tree = out_avals.tree\n      if fun_out_tree != tree_util.tree_structure(None):\n        raise ValueError(\n            f\"The kernel function in mpmd_map {debug_info.func_src_info}\"\n            f\" should return None. It returns a PyTree: {fun_out_tree}.\"\n        )\n      if jaxpr.consts:\n        _error_if_non_ref_consts(jaxpr.consts, debug_info)\n      jaxprs.append(jaxpr)\n      consts_per_fn.append(jaxpr.consts)\n\n    if any(consts_per_fn):\n      # If we close over any constants in the kernel functions, we need to\n      # deduplicate them and then unify the jaxpr signatures.\n      jaxprs, consts = _dedup_consts_and_unify_jaxpr_signatures(\n          jaxprs,\n          consts_per_fn,\n          flat_args,\n          unflat_in_avals,\n          unflat_out_avals,\n          flat_kernel_avals,","sourceCodeStart":968,"sourceCodeEnd":1004,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mpmd.py#L968-L1004","documentation":"A Pallas MPMD kernel function must return None because its outputs are written in-place through Ref arguments. After tracing the function, if the resulting output PyTree structure differs from tree_structure(None) (i.e. the function returned any value), a ValueError is raised.","triggerScenarios":"Passing a kernel to mpmd_map that ends with a return statement (e.g. return out, or implicitly returning the last expression), instead of only mutating Refs and returning nothing.","commonSituations":"Porting a pure-functional Pallas kernel (e.g. from pallas_call / TPU kernels that return outputs) to the MPMD in-place API; kernels written in functional style with an implicit return of the last expression in Python.","solutions":["Remove the return value from the kernel so it returns None; write results into the output Ref arguments","Replace 'return result' with 'out_ref[...] = result' and delete the return","Refactor any implicit return (last bare expression) in the kernel body"],"exampleFix":"# before\ndef kernel(x_ref, o_ref):\n    return x_ref[...] * 2  # returns a value\n# after\ndef kernel(x_ref, o_ref):\n    o_ref[...] = x_ref[...] * 2  # returns None","handlingStrategy":"validation","validationCode":"kernel_return = kernel(*dummy_refs)\nassert kernel_return is None, 'mpmd_map kernel must return None'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Write MPMD kernels that only mutate Refs and never return a value","Add a smoke test that calls the kernel with mock refs and asserts the result is None"],"tags":["jax","pallas","mpmd","kernel-api","inplace-semantics"],"backgroundTag":"kernel-must-return-none","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}