original_docstring: Relative attention score for the positional encodings
docstring: Relative attention score for the positional encodings
docstring_tokens: Relative attention score for the positional encodings
code: def relative_positional_attention(self, position_embeds, q_head, context_len, cls_mask=None):
          """Relative attention score for the positional encodings"""
          # q_head has shape batch_size x sea_len x n_head x d_head
          if self.config.attention_type == "factorized":
              # Notations from the paper, appending A.2.2, final formula (https://arxiv.org/abs/2006.03236)
              # phi and pi have shape seq_len x d_model, psi and omega have shape context_len x d_model
              phi, pi, psi, omega = position_embeds
              # Shape n_head x d_head
              u = self.r_r_bias * self.scale
              # Shape d_model x n_head x d_head
              w_r = self.r_kernel
  
              # Shape batch_size x sea_len x n_head x d_model
              q_r_attention = torch.einsum("binh,dnh->bind", q_head + u, w_r)
              q_r_attention_1 = q_r_attention * phi[:, None]
              q_r_attention_2 = q_r_attention * pi[:, None]
  
              # Shape batch_size x n_head x seq_len x context_len
              positional_attn = torch.einsum("bind,jd->bnij", q_r_attention_1, psi) + torch.einsum(
                  "bind,jd->bnij", q_r_attention_2, omega
              )
          else:
              shift = 2 if q_head.shape[1] != context_len else 1
              # Notations from the paper, appending A.2.1, final formula (https://arxiv.org/abs/2006.03236)
              # Grab the proper positional encoding, shape max_rel_len x d_model
              r = position_embeds[self.block_index][shift - 1]
              # Shape n_head x d_head
              v = self.r_r_bias * self.scale
              # Shape d_model x n_head x d_head
              w_r = self.r_kernel
  
              # Shape max_rel_len x n_head x d_model
              r_head = torch.einsum("td,dnh->tnh", r, w_r)
              # Shape batch_size x n_head x seq_len x max_rel_len
              positional_attn = torch.einsum("binh,tnh->bnit", q_head + v, r_head)
              # Shape batch_size x n_head x seq_len x context_len
              positional_attn = _relative_shift_gather(positional_attn, context_len, shift)
  
          if cls_mask is not None:
              positional_attn *= cls_mask
          return positional_attn
code_tokens: def relative_positional_attention ( self , position_embeds , q_head , context_len , cls_mask = None ) : if self . config . attention_type == "factorized" : phi , pi , psi , omega = position_embeds u = self . r_r_bias * self . scale w_r = self . r_kernel q_r_attention = torch . einsum ( "binh,dnh->bind" , q_head + u , w_r ) q_r_attention_1 = q_r_attention * phi [ : , None ] q_r_attention_2 = q_r_attention * pi [ : , None ] positional_attn = torch . einsum ( "bind,jd->bnij" , q_r_attention_1 , psi ) + torch . einsum ( "bind,jd->bnij" , q_r_attention_2 , omega ) else : shift = 2 if q_head . shape [ 1 ] != context_len else 1 r = position_embeds [ self . block_index ] [ shift - 1 ] v = self . r_r_bias * self . scale w_r = self . r_kernel r_head = torch . einsum ( "td,dnh->tnh" , r , w_r ) positional_attn = torch . einsum ( "binh,tnh->bnit" , q_head + v , r_head ) positional_attn = _relative_shift_gather ( positional_attn , context_len , shift ) if cls_mask is not None : positional_attn *= cls_mask return positional_attn
short_docstring: Relative attention score for the positional encodings
short_docstring_tokens: Relative attention score for the positional encodings
comment: """Relative attention score for the positional encodings"""$SEP$# q_head has shape batch_size x sea_len x n_head x d_head$SEP$# Notations from the paper, appending A.2.2, final formula (https://arxiv.org/abs/2006.03236)$SEP$# phi and pi have shape seq_len x d_model, psi and omega have shape context_len x d_model$SEP$# Shape n_head x d_head$SEP$# Shape d_model x n_head x d_head$SEP$# Shape batch_size x sea_len x n_head x d_model$SEP$# Shape batch_size x n_head x seq_len x context_len$SEP$# Notations from the paper, appending A.2.1, final formula (https://arxiv.org/abs/2006.03236)$SEP$# Grab the proper positional encoding, shape max_rel_len x d_model$SEP$# Shape n_head x d_head$SEP$# Shape d_model x n_head x d_head$SEP$# Shape max_rel_len x n_head x d_model$SEP$# Shape batch_size x n_head x seq_len x max_rel_len$SEP$# Shape batch_size x n_head x seq_len x context_len