code: def _get_partition_name(params: tp.Mapping[tp.Any, tp.Any]) -> str:
      if PARTITION_NAME not in params:
          raise ValueError(
              'Trying to transform a Partitioned variable but "partition_name" '
              f"is not specified in scan_metadata: {params}"
          )
      return params[PARTITION_NAME]
  
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 3 for calling a global function with example usage: _get_partition_name(params) and returns: params