class_docstr: Pooler layer.
  
  Pool hidden states of a specific token (for example start of the
  sequence) and add a linear transformation followed by a tanh.
  
  Arguments:
      hidden_size: hidden size
      init_method: weight initialization method for the linear layer.
          bias is set to zero.
class_docstr_tok: ["Pooler","layer",".","Pool","hidden","states","of","a","specific","token","(","for","example","start","of","the","sequence",")","and","add","a","linear","transformation","followed","by","a","tanh",".","Arguments",":","hidden_size",":","hidden","size","init_method",":","weight","initialization","method","for","the","linear","layer",".","bias","is","set","to","zero","."]
code: def forward(self, hidden_states, sequence_index=0):
      # hidden_states: [s, b, h]
      # sequence_index: index of the token to pool.
  
      # gather data along sequence dimensions
      # same pooler is run on all tensor parallel nodes
      if self.sequence_parallel:
          hidden_states = (
              tensor_parallel.gather_from_sequence_parallel_region(
                  hidden_states, tensor_parallel_output_grad=False
              )
          )
  
      pooled = hidden_states[sequence_index, :, :]
      pooled = self.dense(pooled)
      pooled = torch.tanh(pooled)
      return pooled
  
imports: import torch
  import torch.nn.functional
  from megatron import get_args
  from megatron.core import mpu, tensor_parallel
  from megatron.core.enums import ModelType
  from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
  from .enums import AttnMaskType, LayerType
  from .module import MegatronModule
  from .transformer import ParallelTransformer
  from .utils import get_linear_layer
  from .utils import init_method_normal, scaled_init_method_normal
label_desc: Use image node_id 2 for calling the Pooler obj's underlying member method code with example usage: obj.forward(hidden_states, sequence_index) and returns: pooled