code: def get_partition_spec(tree: A) -> A:
      """Extracts a PartitionSpec tree from a PyTree containing ``Variable`` values."""
  
      def _maybe_replicate(x):
          if hasattr(x, "shape"):
              return PartitionSpec()
          else:
              return None
  
      def f(x):
          if isinstance(x, variables.Variable):
              if isinstance(x, HasSharding) and x.sharding:
                  return x.replace(value=PartitionSpec(*x.sharding))
              else:
                  return x.replace(value=_maybe_replicate(x.value))
  
          return _maybe_replicate(x)
  
      return jax.tree_map(
          f,
          tree,
          is_leaf=lambda x: isinstance(x, variables.Variable)
          and not isinstance(x, TreeNode),
      )
  
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 4 for calling a global function with example usage: get_partition_spec(tree) and returns: jax, _maybe_replicate, PartitionSpec, None, x, x