class_docstr: Computes span representations by generating an unnormalized attention score for each
  word in the document. Spans representations are computed with respect to these
  scores by normalising the attention scores for words inside the span.
  
  Given these attention distributions over every span, this module weights the
  corresponding vector representations of the words in the span by this distribution,
  returning a weighted representation of each span.
  
  Parameters
  ----------
  input_dim : ``int``, required.
      The final dimension of the ``sequence_tensor``.
  
  Returns
  -------
  attended_text_embeddings : ``torch.FloatTensor``.
      A tensor of shape (batch_size, num_spans, input_dim), which each span representation
      is formed by locally normalising a global attention over the sequence. The only way
      in which the attention distribution differs over different spans is in the set of words
      over which they are normalized.
class_docstr_tok: ["Computes","span","representations","by","generating","an","unnormalized","attention","score","for","each","word","in","the","document",".","Spans","representations","are","computed","with","respect","to","these","scores","by","normalising","the","attention","scores","for","words","inside","the","span",".","Given","these","attention","distributions","over","every","span",",","this","module","weights","the","corresponding","vector","representations","of","the","words","in","the","span","by","this","distribution",",","returning","a","weighted","representation","of","each","span",".","Parameters","--","--","--","--","--","input_dim",":","``","int","``",",","required",".","The","final","dimension","of","the","``","sequence_tensor","``",".","Returns","--","--","--","-","attended_text_embeddings",":","``","torch.FloatTensor","``",".","A","tensor","of","shape","(","batch_size",",","num_spans",",","input_dim",")",",","which","each","span","representation","is","formed","by","locally","normalising","a","global","attention","over","the","sequence",".","The","only","way","in","which","the","attention","distribution","differs","over","different","spans","is","in","the","set","of","words","over","which","they","are","normalized","."]
code: def forward(
      self,
      sequence_tensor: torch.FloatTensor,
      span_indices: torch.LongTensor,
      sequence_mask: torch.LongTensor = None,
      span_indices_mask: torch.LongTensor = None,
  ) -> torch.FloatTensor:
      # both of shape (batch_size, num_spans, 1)
      span_starts, span_ends = span_indices.split(1, dim=-1)
  
      # shape (batch_size, num_spans, 1)
      # These span widths are off by 1, because the span ends are `inclusive`.
      span_widths = span_ends - span_starts
  
      # We need to know the maximum span width so we can
      # generate indices to extract the spans from the sequence tensor.
      # These indices will then get masked below, such that if the length
      # of a given span is smaller than the max, the rest of the values
      # are masked.
      max_batch_span_width = span_widths.max().item() + 1
  
      # shape (batch_size, sequence_length, 1)
      global_attention_logits = self._global_attention(sequence_tensor)
  
      # Shape: (1, 1, max_batch_span_width)
      max_span_range_indices = get_range_vector(
          max_batch_span_width, get_device_of(sequence_tensor)
      ).view(1, 1, -1)
      # Shape: (batch_size, num_spans, max_batch_span_width)
      # This is a broadcasted comparison - for each span we are considering,
      # we are creating a range vector of size max_span_width, but masking values
      # which are greater than the actual length of the span.
      #
      # We're using <= here (and for the mask below) because the span ends are
      # inclusive, so we want to include indices which are equal to span_widths rather
      # than using it as a non-inclusive upper bound.
      span_mask = (max_span_range_indices <= span_widths).float()
      raw_span_indices = span_ends - max_span_range_indices
      # We also don't want to include span indices which are less than zero,
      # which happens because some spans near the beginning of the sequence
      # have an end index < max_batch_span_width, so we add this to the mask here.
      span_mask = span_mask * (raw_span_indices >= 0).float()
      span_indices = torch.nn.functional.relu(
          raw_span_indices.float()
      ).long()
  
      # Shape: (batch_size * num_spans * max_batch_span_width)
      flat_span_indices = flatten_and_batch_shift_indices(
          span_indices, sequence_tensor.size(1)
      )
  
      # Shape: (batch_size, num_spans, max_batch_span_width, embedding_dim)
      span_embeddings = batched_index_select(
          sequence_tensor, span_indices, flat_span_indices
      )
  
      # Shape: (batch_size, num_spans, max_batch_span_width)
      span_attention_logits = batched_index_select(
          global_attention_logits, span_indices, flat_span_indices
      ).squeeze(-1)
      # Shape: (batch_size, num_spans, max_batch_span_width)
      span_attention_weights = masked_softmax(
          span_attention_logits, span_mask
      )
  
      # Do a weighted sum of the embedded spans with
      # respect to the normalised attention distributions.
      # Shape: (batch_size, num_spans, embedding_dim)
      attended_text_embeddings = weighted_sum(
          span_embeddings, span_attention_weights
      )
  
      if span_indices_mask is not None:
          # Above we were masking the widths of spans with respect to the max
          # span width in the batch. Here we are masking the spans which were
          # originally passed in as padding.
          return (
              attended_text_embeddings
              * span_indices_mask.unsqueeze(-1).float()
          )
  
      return attended_text_embeddings
  
file_path: jiant/jiant/ext/allennlp.py
imports: from typing import Optional
  import torch
  import torch.nn
label_desc: Use image node_id 4 for calling the SelfAttentiveSpanExtractor obj's underlying member method code with example usage: obj.forward(sequence_tensor, span_indices, sequence_mask, span_indices_mask) and returns: attended_text_embeddings, unknown