code: def get_named_sharding(tree: A, mesh: jax.sharding.Mesh) -> A:
      spec = get_partition_spec(tree)
      sharding = jax.tree_map(
          lambda p: jax.sharding.NamedSharding(mesh, p), spec
      )
      return sharding
  
imports: import functools
  import typing
  import jax
  from jax.experimental import maps
  from jax.sharding import Mesh, PartitionSpec
  from flax.experimental.nnx.nnx import variables
  from flax.experimental.nnx.nnx.pytreelib import TreeNode
  from flax.experimental.nnx.nnx.state import State
label_desc: Use image node_id 5 for calling a global function with example usage: get_named_sharding(tree, mesh) and returns: sharding