class_docstr: Embedding parallelized in the vocabulary dimension.
  
  This is mainly adapted from torch.nn.Embedding and all the default
  values are kept.
  Arguments:
      num_embeddings: vocabulary size.
      embedding_dim: size of hidden state.
      init_method: method to initialize weights.
code: def __init__(
      self,
      num_embeddings: int,
      embedding_dim: int,
      init_method=init.xavier_normal_,
      *,
      params_dtype: torch.dtype = torch.float32,
      use_cpu_initialization: bool = False,
  ):
      super().__init__()
      # Keep the input dimensions.
      self.num_embeddings = num_embeddings
      self.embedding_dim = embedding_dim
      # Set the detauls for compatibility.
      self.padding_idx = None
      self.max_norm = None
      self.norm_type = 2.0
      self.scale_grad_by_freq = False
      self.sparse = False
      self._weight = None
      self.tensor_model_parallel_size = (
          get_tensor_model_parallel_world_size()
      )
      # Divide the weight matrix along the vocabulary dimension.
      (
          self.vocab_start_index,
          self.vocab_end_index,
      ) = VocabUtility.vocab_range_from_global_vocab_size(
          self.num_embeddings,
          get_tensor_model_parallel_rank(),
          self.tensor_model_parallel_size,
      )
      self.num_embeddings_per_partition = (
          self.vocab_end_index - self.vocab_start_index
      )
  
      # Allocate weights and initialize.
      if use_cpu_initialization:
          self.weight = Parameter(
              torch.empty(
                  self.num_embeddings_per_partition,
                  self.embedding_dim,
                  dtype=params_dtype,
              )
          )
          _initialize_affine_weight_cpu(
              self.weight,
              self.num_embeddings,
              self.embedding_dim,
              self.num_embeddings_per_partition,
              0,
              init_method,
              params_dtype=params_dtype,
          )
      else:
          self.weight = Parameter(
              torch.empty(
                  self.num_embeddings_per_partition,
                  self.embedding_dim,
                  device=torch.cuda.current_device(),
                  dtype=params_dtype,
              )
          )
          _initialize_affine_weight_gpu(
              self.weight, init_method, partition_dim=0, stride=1
          )
  
imports: from typing import Optional, Dict, Tuple, List
  import warnings
  import torch
  import torch.nn.functional
  import torch.nn.init
  from torch.nn.parameter import Parameter
  from apex._autocast_utils import _cast_if_autocast_enabled
  from apex.transformer.parallel_state import get_tensor_model_parallel_group
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
  from apex.transformer.utils import divide
  from apex.transformer.tensor_parallel.mappings import copy_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import scatter_to_tensor_model_parallel_region
  from apex.transformer.tensor_parallel.mappings import reduce_scatter_to_sequence_parallel_region
  from apex.transformer.tensor_parallel.random import get_cuda_rng_tracker
  from apex.transformer.tensor_parallel.utils import VocabUtility
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 1 to create a new VocabParallelEmbedding object from inherited base classes: torch.nn with example: obj = VocabParallelEmbedding(num_embeddings, embedding_dim, init_method)