class_docstr: Wrapper for partitioning metadata.
  
  ``Partitioned`` is used to extend variables with partitioning information
  required for ``jax.experimental.pjit``.
  
  The easiest way to define Partitioned variables is by using the
  ``with_partitioning`` wrapper around the variable initializer.
  
  Example::
  
    class MLP(nn.Module):
      hidden_size: int
      @nn.compact
      def __call__(self, x):
        ki = nn.linear.default_kernel_init
        h = nn.Dense(
            self.hidden_size,
            kernel_init=nn.with_partitioning(ki, ('data', 'model')))(x)
        h = nn.relu(h)
        return nn.Dense(
            x.shape[-1],
            kernel_init=nn.with_partitioning(ki, ('model', 'data')))(h)
    mlp = MLP(4096)
    x = jnp.ones((8 * 1024, 1024))
    # use eval_shape to get the Partitioned instances for the variables.
    # this way we can determine the PartitionSpecs for the init variables
    # before we call the init fn.
    var_spec = nn.get_partition_spec(
        jax.eval_shape(mlp.init, random.key(0), x))
    init_fn = mesh(pjit(mlp.init,
                        (None, PartitionSpec("data", "model")), var_spec))
    variables = init_fn(random.key(0), x)
    apply_fn = mesh(pjit(
        mlp.apply,
        (var_spec, PartitionSpec("data", "model")),
         PartitionSpec("data", "model")))
    apply_fn(variables, x)
  
  
  ``Partitioned`` values can gain additional axes when using transformations
  like ``nn.vmap`` and ``nn.scan``. In this case you can specify the name of
  the new axis with the `metadata_params` args in vmap/scan::
  
    class Model(nn.Module):
    @nn.compact
    def __call__(self, x):
      def body(mdl, c):
        c = MLP(4096)(c)
        return c, ()
      c, _ = nn.scan(
          body, variable_axes={"params": 0}, split_rngs={"params": 0}, length=8,
          metadata_params={nn.meta.PARTITION_NAME: "layers"})(self, x)
      return c
code: def _get_partition_name(self, params: Dict[Any, Any]) -> str:
      if PARTITION_NAME not in params:
          raise errors.PartitioningUnspecifiedError(self)
      return params[PARTITION_NAME]
  
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 3 for calling the Partitioned obj's underlying member method code with example usage: obj._get_partition_name(params) and returns: params