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 forward(self, input_):
      if self.tensor_model_parallel_size > 1:
          # Build the mask.
          input_mask = (input_ < self.vocab_start_index) | (
              input_ >= self.vocab_end_index
          )
          # Mask the input.
          masked_input = input_.clone() - self.vocab_start_index
          masked_input[input_mask] = 0
      else:
          masked_input = input_
          # Get the embeddings.
      output_parallel = F.embedding(
          masked_input,
          self.weight,
          self.padding_idx,
          self.max_norm,
          self.norm_type,
          self.scale_grad_by_freq,
          self.sparse,
      )
      # Mask the output embedding.
      if self.tensor_model_parallel_size > 1:
          output_parallel[input_mask, :] = 0.0
      # Reduce across all the model parallel GPUs.
      output = reduce_from_tensor_model_parallel_region(output_parallel)
      return output
  
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 2 for calling the VocabParallelEmbedding obj's underlying member method code with example usage: obj.forward(input_) and returns: output