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: Project hidden [B x N x H] to fixed-size [B x k x H]
  
  return_ortho_loss - if True returns loss term to encourage
                        orthogonal attention vectors
docstr_tok: ["Project","hidden","[","B","x","N","x","H","]","to","fixed-size","[","B","x","k","x","H","]","return_ortho_loss","-","if","True","returns","loss","term","to","encourage","orthogonal","attention","vectors"]
code: def forward(self, hidden, hidden_mask=None, return_ortho_loss=False):
      """
      Project hidden [B x N x H] to fixed-size [B x k x H]
  
      return_ortho_loss - if True returns loss term to encourage
                            orthogonal attention vectors
      """
  
      attention_scores = self.W2(
          self.act(self.W1(hidden) / self.attn_scale) / self.attn_scale
      ).transpose(-1, -2)
  
      attention_mask = form_attention_mask(hidden_mask)
      if attention_mask is not None:
          attention_mask.squeeze_(1)
          attention_scores = attention_scores + attention_mask.to(
              attention_scores.dtype
          )
  
      A = torch.softmax(attention_scores, dim=-1)
      M = A @ hidden
  
      if return_ortho_loss:
          ortho_loss = (
              ((A @ A.transpose(-1, -2)) - torch.eye(self.k).type_as(A))
              .pow(2)
              .sum()
          )
  
          return M, ortho_loss
      else:
          return M
  
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 2 for calling the AttentionBridge obj's underlying member method code with example usage: obj.forward(hidden, hidden_mask, return_ortho_loss) and returns: M,  ortho_loss, M