code: def get_axis_names(axes_metadata):
      """Gets axis names for variables as logical PartitionSpecs.
  
      Args:
        axes_metadata: a single axes-metadata collection from a flax-initialized
          set of collections.
  
      Returns:
        Collection of Partitionspecs with logical axis names, with the "_axes"
        suffix on variable names removed to match original variable collection for
        annotations.
      """
  
      def leaf_rewrite(x):
          return None if x is None else jax.sharding.PartitionSpec(*x)
  
      def rewrite(tree):
          return jax.tree_util.tree_map(
              leaf_rewrite, tree, is_leaf=_is_logical_spec
          )
  
      axes_metadata = unfreeze(
          axes_metadata
      )  # pytype: disable=wrong-arg-types
      flat_dict = {
          re.sub(r"_axes$", "", "/".join(k)): rewrite(v.names)
          for k, v in flatten_dict(axes_metadata).items()
      }
      return freeze(
          unflatten_dict(
              {tuple(k.split("/")): v for k, v in flat_dict.items()}
          )
      )
  
imports: import functools
  import re
  from typing import Any, Callable, Mapping, Optional, Tuple
  import flax
  from flax import linen
  from flax import struct
  from flax.core.frozen_dict import freeze
  from flax.core.frozen_dict import unfreeze
  from flax.core.lift import In
  from flax.core.lift import Out
  from flax.linen.spmd import _axis_rules
  from flax.linen.spmd import _AxisRules
  from flax.linen.spmd import _is_logical_spec
  from flax.linen.spmd import _with_sharding_constraint
  from flax.linen.spmd import Array
  from flax.linen.spmd import ArrayPytree
  from flax.linen.spmd import get_logical_axis_rules
  from flax.linen.spmd import logical_axis_rules
  from flax.linen.spmd import logical_to_mesh
  from flax.linen.spmd import logical_to_mesh_axes
  from flax.linen.spmd import LogicalPartitionSpec
  from flax.linen.spmd import LogicalPartitionSpecPytree
  from flax.linen.spmd import LogicalRules
  from flax.linen.spmd import PartitionSpecPytree
  from flax.linen.spmd import RulesFallback
  from flax.linen.spmd import set_logical_axis_rules
  from flax.linen.spmd import with_logical_constraint
  from flax.traverse_util import flatten_dict
  from flax.traverse_util import unflatten_dict
  import jax
label_desc: Use image node_id 5 for calling a global function with example usage: get_axis_names(axes_metadata) and returns: freeze, unknown, jax