code: def get_partition_spec(tree: Any) -> Any:
      """Extracts a PartitionSpec tree from a PyTree containing ``Partitioned`` values."""
  
      def f(x):
          if isinstance(x, Partitioned):
              return x.get_partition_spec()
          # Unboxed arrays, which should be replicated across all devices
          elif hasattr(x, "shape"):
              return jax.sharding.PartitionSpec()
          else:
              return None
  
      return jax.tree_map(
          f, tree, is_leaf=lambda x: isinstance(x, Partitioned)
      )
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Generic, Optional, Tuple, TypeVar, Union
  import jax
  from jax.experimental import maps
  from flax import errors, struct
label_desc: Use image node_id 9 for calling a global function with example usage: get_partition_spec(tree) and returns: jax, x, jax, None