jax-ml/jax · error · ValueError

The kernel function in mpmd_map {debug_info.func_src_info} s

Error message

The kernel function in mpmd_map {debug_info.func_src_info} should return None. It returns a PyTree: {fun_out_tree}.

What it means

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.

Source

Thrown at jax/_src/pallas/mpmd.py:986

    consts_per_fn = []
    debug_infos = [api_util.debug_info("mpmd_map", fn, kernel_arg_avals, kernel_kwarg_avals)
                   for _, fn in meshes_and_fns]
    if name is not None:
      debug_infos = [di.replace_func_name(name) for di in debug_infos]
    # If names are non-distinct (e.g. because user passed multiple functions
    # with the same name, or because of the name= arg handled just above),
    # uniquify them with the core type.
    if len({di.func_name for di in debug_infos}) != len(debug_infos):
      debug_infos = [di.replace_func_name(f"{di.func_name}__{mesh.core_type}")
                     for di, mesh in zip(debug_infos, meshes)]
    for (mesh, fn), debug_info in zip(meshes_and_fns, debug_infos):
      with mpmd_map_tracing_context(mesh, all_meshes):
        jaxpr, out_avals = pe.trace_to_jaxpr(
            fn, in_avals_ft, debug_info
        )
      fun_out_tree = out_avals.tree
      if fun_out_tree != tree_util.tree_structure(None):
        raise ValueError(
            f"The kernel function in mpmd_map {debug_info.func_src_info}"
            f" should return None. It returns a PyTree: {fun_out_tree}."
        )
      if jaxpr.consts:
        _error_if_non_ref_consts(jaxpr.consts, debug_info)
      jaxprs.append(jaxpr)
      consts_per_fn.append(jaxpr.consts)

    if any(consts_per_fn):
      # If we close over any constants in the kernel functions, we need to
      # deduplicate them and then unify the jaxpr signatures.
      jaxprs, consts = _dedup_consts_and_unify_jaxpr_signatures(
          jaxprs,
          consts_per_fn,
          flat_args,
          unflat_in_avals,
          unflat_out_avals,
          flat_kernel_avals,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the return value from the kernel so it returns None; write results into the output Ref arguments
2. Replace 'return result' with 'out_ref[...] = result' and delete the return
3. Refactor any implicit return (last bare expression) in the kernel body

Example fix

# before
def kernel(x_ref, o_ref):
    return x_ref[...] * 2  # returns a value
# after
def kernel(x_ref, o_ref):
    o_ref[...] = x_ref[...] * 2  # returns None

kernel_return = kernel(*dummy_refs)
assert kernel_return is None, 'mpmd_map kernel must return None'

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Cannot pass the same ref into a mpmd map multiple times
• MPMD map with more than one mesh requires scratch_type to ha
• This gmm kernel only supports either (m, k) x (g, k, n) -> (
• Group sizes {group_sizes.shape=} must match first dimension
• Explicit sharding is not currently supported in the pallas-t

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