docstr: Calculate attention for the given key, query value
  
  Where B is the number of items stacked together, N is the length:
  The key should be BxNxD
  The query is BxD
  The value is BxNxD
  
  If mask is specified, it should be BxN of True/False values,
  where True means that location is masked out
  
  Reshapes and reorders are used to handle num_heads
  
  Return will be softmax(query x key^T) * value
  of size BxD
code: def attention(self, key, query, value, mask=None):
      """
      Calculate attention for the given key, query value
  
      Where B is the number of items stacked together, N is the length:
      The key should be BxNxD
      The query is BxD
      The value is BxNxD
  
      If mask is specified, it should be BxN of True/False values,
      where True means that location is masked out
  
      Reshapes and reorders are used to handle num_heads
  
      Return will be softmax(query x key^T) * value
      of size BxD
      """
      B = key.shape[0]
      N = key.shape[1]
      D = key.shape[2]
  
      H = self.num_heads
  
      # query is now BxDx1
      query = query.unsqueeze(2)
      # BxHxD/Hx1
      query = query.reshape((B, H, -1, 1))
  
      # BxNxHxD/H
      key = key.reshape((B, N, H, -1))
      # BxHxNxD/H
      key = key.transpose(1, 2)
  
      # BxNxHxD/H
      value = value.reshape((B, N, H, -1))
      # BxHxNxD/H
      value = value.transpose(1, 2)
  
      # BxHxNxD/H x BxHxD/Hx1
      # result shape: BxHxN
      attn = (
          torch.matmul(key, query).squeeze(3)
          * self.inv_sqrt_output_size
      )
      if mask is not None:
          # mask goes from BxN -> Bx1xN
          mask = mask.unsqueeze(1)
          mask = mask.expand(-1, H, -1)
          attn.masked_fill_(mask, float("-inf"))
      # attn shape will now be BxHx1xN
      attn = torch.softmax(attn, dim=2).unsqueeze(2)
      # BxHx1xN x BxHxNxD/H -> BxHxD/H
      output = torch.matmul(attn, value).squeeze(2)
      output = output.reshape(B, -1)
      return output
  
file_path: stanza/stanza/models/constituency/transformer_tree_stack.py
imports: from collections import namedtuple
  import torch
  import torch.nn
  from stanza.models.constituency.positional_encoding import SinusoidalEncoding
  from stanza.models.constituency.tree_stack import TreeStack
label_desc: Use image node_id 2 for calling the TransformerTreeStack obj's underlying member method code with example usage: obj.attention(key, query, value, mask) and returns: output