code: def flatten_axes(
      name, treedef, axis_tree, *, kws=False, tupled_args=False
  ):
      # given an axis spec tree axis_tree (a pytree with integers and Nones at the
      # leaves, i.e. the Nones are to be considered leaves) that is a tree prefix of
      # the given treedef, build a complete axis spec tree with the same structure
      # and return the flattened result
      # TODO(mattjj,phawkins): improve this implementation
      proxy = object()
      dummy = tree_unflatten(treedef, [object()] * treedef.num_leaves)
      axes = []
      add_leaves = lambda i, x: axes.extend(
          [i] * len(tree_flatten(x)[0])
      )
      try:
          tree_map(add_leaves, _replace_nones(proxy, axis_tree), dummy)
      except ValueError:
          if kws:
              # if keyword arguments are included in the tree, we make adapt the error
              # message only to be about the positional arguments
              treedef, _ = treedef_children(treedef)
              axis_tree, _ = axis_tree
          hint = ""
          if tupled_args:
              hint += (
                  f" Note that {name} that are non-trivial pytrees should always be "
                  f"wrapped in a tuple representing the argument list."
              )
              if len(treedef.children()) == 1:
                  try:
                      flatten_axes(name, treedef, (axis_tree,))
                  except ValueError:
                      pass  # That's not the issue.
                  else:
                      hint += (
                          f" In particular, you're passing in a single argument which "
                          f"means that {name} might need to be wrapped in "
                          f"a singleton tuple."
                      )
          raise ValueError(
              f"{name} specification must be a tree prefix of the "
              f"corresponding value, got specification {axis_tree} "
              f"for value tree {treedef}.{hint}"
          ) from None
      axes = [None if a is proxy else a for a in axes]
      assert len(axes) == treedef.num_leaves
      return axes
  
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 22 for calling a global function with example usage: flatten_axes(name, treedef, axis_tree) and returns: axes