code: def _flat_out_axes(leaves, treedef, *args, **kwargs):
      ans = yield args, kwargs
      spec = tree_unflatten(treedef, leaves)
      try:
          spec_flat = tuple(
              broadcast_prefix(spec, ans, is_leaf=lambda x: x is None)
          )
      except ValueError:
          e, *_ = prefix_errors(spec, ans)
          # TODO(mattjj): currently hardcoded for pmap; generalize to vmap in followup
          (msg,) = e("pmap out_axes").args
          msg += (
              "\n\nThe full pytree is the output of the pmapped function. Ensure "
              "that the `out_axes` argument to `pmap` is a pytree prefix of the "
              "pmapped function's output."
          )
          raise ValueError(msg) from None
      yield ans, spec_flat
  
imports: from __future__ import annotations
  from collections.abc import Iterable, Sequence
  import inspect
  import operator
  from functools import partial
  from typing import Any, Callable
  import warnings
  import numpy
  from jax._src import core
  from jax._src import dtypes
  from jax._src.abstract_arrays import numpy_scalar_types
  from jax._src.core import ShapedArray
  from jax._src.tree_util import PyTreeDef, tree_flatten, tree_unflatten, tree_map, tree_structure, treedef_children, generate_key_paths, keystr, broadcast_prefix, prefix_errors
  from jax._src.tree_util import _replace_nones
  from jax._src import linear_util
  from jax._src.linear_util import TracingDebugInfo
  from jax._src.util import safe_map, WrapKwArgs, Hashable, HashableFunction, Unhashable
  from jax._src import traceback_util
label_desc: Use image node_id 24 for calling a global function with example usage: _flat_out_axes(leaves, treedef) without return types