text: from typing import Optional, Tuple
  import torch
  import torch.nn.functional as F
  from torch.overrides import has_torch_function
  from torch import Tensor, softmax
  
  def _NN_attention(
      q: Tensor,
      k: Tensor,
      v: Tensor,
      NN,
      attn_mask: Optional[Tensor] = None,
      dropout_p: float = 0.0,
  ) -> Tuple[Tensor, Tensor]:
      q1 = q.unsqueeze(1)
      k1 = k.unsqueeze(2)
      q2 = q1.expand(-1, k.shape[1], -1, -1)
      k2 = k1.expand(-1, -1, q.shape[1], -1)
      nninp = torch.cat((q2, k2), dim=-1).transpose(-2, -3)
      nninp2 = torch.cat((k2, q2), dim=-1).transpose(-2, -3)
  
      attn = NN(nninp).squeeze(-1)
      attn2 = NN(nninp2).squeeze(-1)
      attn = attn + attn2
      
      if attn_mask is not None:
          attn += attn_mask
      attn = softmax(attn, dim=-1)
      if dropout_p > 0.0:
          attn = F.dropout(attn, p=dropout_p)
      # (B, Nt, Ns) x (B, Ns, E) -> (B, Nt, E)
      output = torch.bmm(attn, v)
      return output, attn
  
  def _in_projection(
      q: Tensor,
      k: Tensor,
      v: Tensor,
      w_q: Tensor,
      w_k: Tensor,
      w_v: Tensor,
      b_q: Optional[Tensor] = None,
      b_k: Optional[Tensor] = None,
      b_v: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Tensor, Tensor]:
      Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
      assert w_q.shape == (Eq, Eq), f"expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}"
      assert w_k.shape == (Eq, Ek), f"expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}"
      assert w_v.shape == (Eq, Ev), f"expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}"
      assert b_q is None or b_q.shape == (Eq,), f"expecting query bias shape of {(Eq,)}, but got {b_q.shape}"
      assert b_k is None or b_k.shape == (Eq,), f"expecting key bias shape of {(Eq,)}, but got {b_k.shape}"
      assert b_v is None or b_v.shape == (Eq,), f"expecting value bias shape of {(Eq,)}, but got {b_v.shape}"
      return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)
  
  def multi_head_attention_forward(
      query: Tensor,
      key: Tensor,
      value: Tensor,
      NN,
      embed_dim_to_check: int,
      num_heads: int,
      in_proj_weight: Tensor,
      in_proj_bias: Optional[Tensor],
      bias_k: Optional[Tensor],
      bias_v: Optional[Tensor],
      add_zero_attn: bool,
      dropout_p: float,
      out_proj_weight: Tensor,
      out_proj_bias: Optional[Tensor],
      training: bool = True,
      key_padding_mask: Optional[Tensor] = None,
      need_weights: bool = True,
      attn_mask: Optional[Tensor] = None,
      use_separate_proj_weight: bool = False,
      q_proj_weight: Optional[Tensor] = None,
      k_proj_weight: Optional[Tensor] = None,
      v_proj_weight: Optional[Tensor] = None,
      static_k: Optional[Tensor] = None,
      static_v: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)
      if has_torch_function(tens_ops):
          assert False
          return handle_torch_function(
              multi_head_attention_forward,
              tens_ops,
              query,
              key,
              value,
              embed_dim_to_check,
              num_heads,
              in_proj_weight,
              in_proj_bias,
              bias_k,
              bias_v,
              add_zero_attn,
              dropout_p,
              out_proj_weight,
              out_proj_bias,
              training=training,
              key_padding_mask=key_padding_mask,
              need_weights=need_weights,
              attn_mask=attn_mask,
              use_separate_proj_weight=use_separate_proj_weight,
              q_proj_weight=q_proj_weight,
              k_proj_weight=k_proj_weight,
              v_proj_weight=v_proj_weight,
              static_k=static_k,
              static_v=static_v,
          )
          
      assert NN is not None
  
      # set up shape vars
      tgt_len, bsz, embed_dim = query.shape
      src_len, _, _ = key.shape
      assert embed_dim == embed_dim_to_check, \
          f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
      if isinstance(embed_dim, torch.Tensor):
          # embed_dim can be a tensor when JIT tracing
          head_dim = embed_dim.div(num_heads, rounding_mode='trunc')
      else:
          head_dim = embed_dim // num_heads
      assert head_dim * num_heads == embed_dim, f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
      if use_separate_proj_weight:
          # allow MHA to have different embedding dimensions when separate projection weights are used
          assert key.shape[:2] == value.shape[:2], \
              f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}"
      else:
          assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
  
      #
      # compute in-projection
      #
      if not use_separate_proj_weight:
          q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
      else:
          assert q_proj_weight is not None, "use_separate_proj_weight is True but q_proj_weight is None"
          assert k_proj_weight is not None, "use_separate_proj_weight is True but k_proj_weight is None"
          assert v_proj_weight is not None, "use_separate_proj_weight is True but v_proj_weight is None"
          if in_proj_bias is None:
              b_q = b_k = b_v = None
          else:
              b_q, b_k, b_v = in_proj_bias.chunk(3)
          q, k, v = _in_projection(query, key, value, q_proj_weight, k_proj_weight, v_proj_weight, b_q, b_k, b_v)
  
      # prep attention mask
      if attn_mask is not None:
          if attn_mask.dtype == torch.uint8:
              warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
              attn_mask = attn_mask.to(torch.bool)
          else:
              assert attn_mask.is_floating_point() or attn_mask.dtype == torch.bool, \
                  f"Only float, byte, and bool types are supported for attn_mask, not {attn_mask.dtype}"
          # ensure attn_mask's dim is 3
          if attn_mask.dim() == 2:
              correct_2d_size = (tgt_len, src_len)
              if attn_mask.shape != correct_2d_size:
                  raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
              attn_mask = attn_mask.unsqueeze(0)
          elif attn_mask.dim() == 3:
              correct_3d_size = (bsz * num_heads, tgt_len, src_len)
              if attn_mask.shape != correct_3d_size:
                  raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
          else:
              raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
  
      # prep key padding mask
      if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
          warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
          key_padding_mask = key_padding_mask.to(torch.bool)
  
      # add bias along batch dimension (currently second)
      if bias_k is not None and bias_v is not None:
          assert static_k is None, "bias cannot be added to static key."
          assert static_v is None, "bias cannot be added to static value."
          k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
          v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
          if attn_mask is not None:
              attn_mask = pad(attn_mask, (0, 1))
          if key_padding_mask is not None:
              key_padding_mask = pad(key_padding_mask, (0, 1))
      else:
          assert bias_k is None
          assert bias_v is None
  
      #
      # reshape q, k, v for multihead attention and make em batch first
      #
      q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
      if static_k is None:
          k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
      else:
          assert static_k.size(0) == bsz * num_heads, \
              f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
          assert static_k.size(2) == head_dim, \
              f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
          k = static_k
      if static_v is None:
          v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
      else:
          assert static_v.size(0) == bsz * num_heads, \
              f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
          assert static_v.size(2) == head_dim, \
              f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
          v = static_v
  
      # add zero attention along batch dimension (now first)
      if add_zero_attn:
          zero_attn_shape = (bsz * num_heads, 1, head_dim)
          k = torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1)
          v = torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1)
          if attn_mask is not None:
              attn_mask = pad(attn_mask, (0, 1))
          if key_padding_mask is not None:
              key_padding_mask = pad(key_padding_mask, (0, 1))
  
      # update source sequence length after adjustments
      src_len = k.size(1)
  
      # merge key padding and attention masks
      if key_padding_mask is not None:
          assert key_padding_mask.shape == (bsz, src_len), \
              f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
          key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).   \
              expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, src_len)
          if attn_mask is None:
              attn_mask = key_padding_mask
          elif attn_mask.dtype == torch.bool:
              attn_mask = attn_mask.logical_or(key_padding_mask)
          else:
              attn_mask = attn_mask.masked_fill(key_padding_mask, float("-inf"))
  
      # convert mask to float
      if attn_mask is not None and attn_mask.dtype == torch.bool:
          new_attn_mask = torch.zeros_like(attn_mask, dtype=torch.float)
          new_attn_mask.masked_fill_(attn_mask, float("-inf"))
          attn_mask = new_attn_mask
  
      # adjust dropout probability
      if not training:
          dropout_p = 0.0
  
      #
      # (deep breath) calculate attention and out projection
      #
      attn_output, attn_output_weights = _NN_attention(q, k, v, NN, attn_mask, dropout_p)
      attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
      attn_output = F.linear(attn_output, out_proj_weight, out_proj_bias)
  
      if need_weights:
          # average attention weights over heads
          attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
          return attn_output, attn_output_weights.sum(dim=1) / num_heads
      else:
          return attn_output, None
  
  
  class CustomMultiheadAttention(torch.nn.MultiheadAttention):
      def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
          super(CustomMultiheadAttention, self).__init__(embed_dim, num_heads, dropout=dropout, bias=bias, add_bias_kv=add_bias_kv, add_zero_attn=add_zero_attn, kdim=kdim, vdim=vdim)
  
      def forward(self, query: Tensor, key: Tensor, value: Tensor, NN, key_padding_mask: Optional[Tensor] = None,
                  need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
  
          if not self._qkv_same_embed_dim:
              return multi_head_attention_forward(
                  query, key, value, NN, self.embed_dim, self.num_heads,
                  self.in_proj_weight, self.in_proj_bias,
                  self.bias_k, self.bias_v, self.add_zero_attn,
                  self.dropout, self.out_proj.weight, self.out_proj.bias,
                  training=self.training,
                  key_padding_mask=key_padding_mask, need_weights=need_weights,
                  attn_mask=attn_mask, use_separate_proj_weight=True,
                  q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
                  v_proj_weight=self.v_proj_weight)
          else:
              return multi_head_attention_forward(
                  query, key, value, NN, self.embed_dim, self.num_heads,
                  self.in_proj_weight, self.in_proj_bias,
                  self.bias_k, self.bias_v, self.add_zero_attn,
                  self.dropout, self.out_proj.weight, self.out_proj.bias,
                  training=self.training,
                  key_padding_mask=key_padding_mask, need_weights=need_weights,
                  attn_mask=attn_mask)
  
  def hack_attention(layer):
      layer._qkv_same_embed_dim = False
      layer.register_parameter('in_proj_weight', None)
      embed_dim = layer.embed_dim
      layer.q_proj_weight = torch.nn.Parameter(torch.Tensor(embed_dim, embed_dim))
      layer.k_proj_weight = torch.nn.Parameter(torch.Tensor(embed_dim, embed_dim))
      layer.v_proj_weight = torch.nn.Parameter(torch.Tensor(embed_dim, embed_dim))
      layer._reset_parameters()
      layer.k_proj_weight = layer.q_proj_weight
      layer.v_proj_weight = layer.q_proj_weight
  
  class CustomTransformerDecoderLayer(torch.nn.TransformerDecoderLayer):
      def __init__(self, d_model, nhead, NN=None, dim_feedforward=2048, dropout=0.1, activation="relu"):
          super(CustomTransformerDecoderLayer, self).__init__(d_model, nhead, dim_feedforward=dim_feedforward, dropout=dropout, activation=activation)
          self.NN = NN
          self.self_attn = CustomMultiheadAttention(d_model, nhead, dropout=dropout)
          self.multihead_attn = CustomMultiheadAttention(d_model, nhead, dropout=dropout)
          hack_attention(self.self_attn)
          hack_attention(self.multihead_attn)
  
      def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None,
                  tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
          r"""Pass the inputs (and mask) through the decoder layer.
  
          Args:
              tgt: the sequence to the decoder layer (required).
              memory: the sequence from the last layer of the encoder (required).
              tgt_mask: the mask for the tgt sequence (optional).
              memory_mask: the mask for the memory sequence (optional).
              tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
              memory_key_padding_mask: the mask for the memory keys per batch (optional).
  
          Shape:
              see the docs in Transformer class.
          """
          tgt2 = self.self_attn(tgt, tgt, tgt, self.NN, attn_mask=tgt_mask,
                                key_padding_mask=tgt_key_padding_mask)[0]
          tgt = tgt + self.dropout1(tgt2)
          tgt = self.norm1(tgt)
          tgt2, ao = self.multihead_attn(tgt, memory, memory, self.NN, attn_mask=memory_mask,
                                     key_padding_mask=memory_key_padding_mask)
          self.attention = ao
          tgt = 1.5*tgt + 0.5*self.dropout2(tgt2)
          tgt = self.norm2(tgt)
          tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
          tgt = 1.5*tgt + 0.5*self.dropout3(tgt2)
          tgt = self.norm3(tgt)
          return tgt
  
  class CustomTransformerEncoderLayer(torch.nn.TransformerEncoderLayer):
      def __init__(self, d_model, nhead, NN=None, dim_feedforward=2048, dropout=0.1, activation="relu"):
          super(CustomTransformerEncoderLayer, self).__init__(d_model, nhead, dim_feedforward=dim_feedforward, dropout=dropout, activation=activation)
          self.NN = NN
          self.self_attn = CustomMultiheadAttention(d_model, nhead, dropout=dropout)
          hack_attention(self.self_attn)
  
      def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
          r"""Pass the input through the encoder layer.
  
          Args:
              src: the sequence to the encoder layer (required).
              src_mask: the mask for the src sequence (optional).
              src_key_padding_mask: the mask for the src keys per batch (optional).
  
          Shape:
              see the docs in Transformer class.
          """
  
          src2 = self.self_attn(src, src, src, self.NN, attn_mask=src_mask,
                                key_padding_mask=src_key_padding_mask)[0]
          src = 1.5*src + 0.5*self.dropout1(src2)
          src = self.norm1(src)
          src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
          src = 1.5*src + 0.5*self.dropout2(src2)
          src = self.norm2(src)
          return src
  