code: def get_sharding(tree: Any, mesh: jax.sharding.Mesh) -> Any:
      """Extracts a jax.sharding tree from a PyTree containing ``Partitioned`` values and a mesh."""
      pspec_tree = get_partition_spec(tree)
      return jax.tree_map(
          lambda x: jax.sharding.NamedSharding(mesh, x), pspec_tree
      )
  
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 10 for calling a global function with example usage: get_sharding(tree, mesh) and returns: jax