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.
  
  Keyword Arguments:
      config: A megatron.core.ModelParallelConfig object
code: def __init__(
      self,
      num_embeddings: int,
      embedding_dim: int,
      *,
      init_method: Callable,
      config: ModelParallelConfig,
  ):
      super(VocabParallelEmbedding, self).__init__()
      # Keep the input dimensions.
      self.num_embeddings = num_embeddings
      self.embedding_dim = embedding_dim
      self.tensor_model_parallel_size = (
          get_tensor_model_parallel_world_size()
      )
      # Divide the weight matrix along the vocaburaly 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 config.use_cpu_initialization:
          self.weight = Parameter(
              torch.empty(
                  self.num_embeddings_per_partition,
                  self.embedding_dim,
                  dtype=config.params_dtype,
              )
          )
          if config.perform_initialization:
              _initialize_affine_weight_cpu(
                  self.weight,
                  self.num_embeddings,
                  self.embedding_dim,
                  self.num_embeddings_per_partition,
                  0,
                  init_method,
                  params_dtype=config.params_dtype,
              )
      else:
          self.weight = Parameter(
              torch.empty(
                  self.num_embeddings_per_partition,
                  self.embedding_dim,
                  device=torch.cuda.current_device(),
                  dtype=config.params_dtype,
              )
          )
          if config.perform_initialization:
              _initialize_affine_weight_gpu(
                  self.weight, init_method, partition_dim=0, stride=1
              )
  
imports: import math
  import os
  import warnings
  from typing import Callable, Optional
  import torch
  import torch.nn.functional
  import torch.nn.init
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.nn.parameter import Parameter
  from megatron.core.model_parallel_config import ModelParallelConfig
  from megatron.core.parallel_state import get_global_memory_buffer, get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
  from .mappings import copy_to_tensor_model_parallel_region, gather_from_sequence_parallel_region, gather_from_tensor_model_parallel_region, reduce_from_tensor_model_parallel_region, reduce_scatter_to_sequence_parallel_region, scatter_to_tensor_model_parallel_region
  from .random import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name
  from .utils import VocabUtility, divide, split_tensor_along_last_dim
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)