docstr: bert_output
      (bs, seqlen, embed_dim)
  mention_bounds: both bounds are inclusive [start, end]
      (bs, num_spans, 2)
docstr_tok: ["bert_output","(","bs",",","seqlen",",","embed_dim",")","mention_bounds",":","both","bounds","are","inclusive","[","start",",","end","]","(","bs",",","num_spans",",","2",")"]
code: def forward(self, bert_output, mention_bounds):
      """
      bert_output
          (bs, seqlen, embed_dim)
      mention_bounds: both bounds are inclusive [start, end]
          (bs, num_spans, 2)
      """
      # get embedding of [CLS] token
      if mention_bounds.size(0) == 0:
          return mention_bounds
      if self.tokens_to_aggregate == "all":
          (
              embedding_ctxt,  # (batch_size, num_spans, max_batch_span_width, embedding_size)
              mask,  # (batch_size, num_spans, max_batch_span_width)
          ) = batched_span_select(
              bert_output,  # (batch_size, sequence_length, embedding_size)
              mention_bounds,  # (batch_size, num_spans, 2)
          )
          embedding_ctxt[~mask] = 0  # 0 out masked elements
          # embedding_ctxt = (batch_size, num_spans, max_batch_span_width, embedding_size)
          if self.aggregate_method.startswith("avg"):
              embedding_ctxt = embedding_ctxt.sum(2) / mask.sum(
                  2
              ).float().unsqueeze(-1)
              # embedding_ctxt = (batch_size, num_spans, embedding_size)
          if self.aggregate_method == "avg_linear":
              embedding_ctxt = self.mention_agg_linear(embedding_ctxt)
              # embedding_ctxt = (batch_size, num_spans, output_dim)
      elif self.tokens_to_aggregate == "fl":
          start_embeddings = batched_index_select(
              bert_output, mention_bounds[:, :, 0]
          )
          end_embeddings = batched_index_select(
              bert_output, mention_bounds[:, :, 1]
          )
          embedding_ctxt = torch.cat(
              [
                  start_embeddings.unsqueeze(2),
                  end_embeddings.unsqueeze(2),
              ],
              dim=2,
          )
          # embedding_ctxt = (batch_size, num_spans, 2, embedding_size)
          if self.aggregate_method == "avg":
              embedding_ctxt = embedding_ctxt.mean(2)
              # embedding_ctxt = (batch_size, num_spans, embedding_size)
          elif self.aggregate_method == "linear":
              embedding_ctxt = embedding_ctxt.view(
                  embedding_ctxt.size(0), embedding_ctxt.size(1), -1
              )
              # embedding_ctxt = (batch_size, num_spans, 2 * embedding_size)
              embedding_ctxt = self.mention_agg_linear(embedding_ctxt)
              # embedding_ctxt = (batch_size, num_spans, output_dim)
      else:
          raise NotImplementedError()
  
      return embedding_ctxt
  
imports: import os
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from tqdm import tqdm
  from collections import OrderedDict
  from pytorch_transformers.modeling_bert import BertPreTrainedModel, BertConfig, BertModel
  from pytorch_transformers.tokenization_bert import BertTokenizer
  from elq.common.ranker_base import BertEncoder, get_model_obj
  from blink.common.optimizer import get_bert_optimizer
  from elq.biencoder.allennlp_span_utils import batched_span_select, batched_index_select
  from elq.biencoder.utils import batch_reshape_mask_left
label_desc: Use image node_id 2 for calling the GetContextEmbedsHead obj's underlying member method code with example usage: obj.forward(bert_output, mention_bounds) and returns: embedding_ctxt, mention_bounds