class_docstr: A multi-head attention bridge to project a variable-size hidden states
  to k hidden states (per attention head).
  
  Code is based on the paper https://arxiv.org/pdf/1703.03130.pdf
class_docstr_tok: ["A","multi-head","attention","bridge","to","project","a","variable-size","hidden","states","to","k","hidden","states","(","per","attention","head",")",".","Code","is","based","on","the","paper","https",":","\/\/arxiv.org\/pdf\/1703.03130.pdf"]
docstr: hidden_size - size of input hidden state
  k - number of attention heads
  bridge_size - size of internal feed forward weights (i.e., attention head size)
docstr_tok: ["hidden_size","-","size","of","input","hidden","state","k","-","number","of","attention","heads","bridge_size","-","size","of","internal","feed","forward","weights","(","i.e.",",","attention","head","size",")"]
code: def __init__(self, hidden_size, k, bridge_size):
      """
      hidden_size - size of input hidden state
      k - number of attention heads
      bridge_size - size of internal feed forward weights (i.e., attention head size)
      """
      super().__init__()
  
      self.hidden_size = hidden_size
      self.k = k
      self.bridge_size = bridge_size
  
      self.attn_scale = np.sqrt(np.sqrt(self.bridge_size))
  
      # build model
  
      self.W1 = torch.nn.Linear(hidden_size, bridge_size, bias=False)
      self.W2 = torch.nn.Linear(bridge_size, k, bias=False)
      self.act = torch.nn.ReLU()
  
imports: import math
  import numpy
  import torch
  from torch import nn
  from torch.nn.functional import gelu
  from nemo.collections.common.parts import form_attention_mask
  from nemo.utils import logging
label_desc: Use image node_id 1 to create a new AttentionBridge object from inherited base classes: torch.nn with example: obj = AttentionBridge(hidden_size, k, bridge_size)