docstr: A variety of pooling options. Some ignore the cls token. Input needs to be B S H.
docstr_tok: ["A","variety","of","pooling","options",".","Some","ignore","the","cls","token",".","Input","needs","to","be","B","S","H","."]
code: def forward(self, hidden_states):
      """A variety of pooling options. Some ignore the cls token. Input needs to be B S H."""
      if self.pool_scheme == "zero_index":
          first_token_tensor = hidden_states[:, 0]
      elif self.pool_scheme == "avg":
          first_token_tensor = hidden_states.mean(dim=1)
      elif self.pool_scheme == "max":
          first_token_tensor = hidden_states.max(dim=1)[0]
      elif self.pool_scheme == "lse":
          first_token_tensor = hidden_states.logsumexp(dim=1)
      else:
          raise ValueError(
              f"Invalid pooling scheme {self.pool_scheme} given."
          )
  
      pooled_output = self.activation(self.dense(first_token_tensor))
      return self.dropout(pooled_output)
  
imports: import torch
  import math
  from typing import Optional, Tuple
  from functools import partial
  from .embeddings import SinusoidalPositional, LearnablePositional, ScaledSinosoidal
  from .attention import get_attention_mechanism
label_desc: Use image node_id 2 for calling the PoolingComponent obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: self