code: # Copyright (c) 2017-present, Facebook, Inc.
  # All rights reserved.
  #
  # This source code is licensed under the license found in the LICENSE file in
  # the root directory of this source tree. An additional grant of patent rights
  # can be found in the PATENTS file in the same directory.
  
  import math
  
  import torch
  from torch import nn
  from torch.nn import Parameter
  import torch.nn.functional as F
  
  from fairseq import utils
  from fairseq.modules.learned_positional_embedding import LearnedPositionalEmbedding
  
  # Adapted from faiserq/modules/multihead_attention to deal with local attention
  # Local attetion masking in combination with padding masking can lead to
  # all -Inf attention rows. This version detects and corrects this situation
  class ProtectedMultiheadAttention(nn.Module):
      """Multi-headed attention.
  
      See "Attention Is All You Need" for more details.
      """
  
      def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False):
          super().__init__()
          self.embed_dim = embed_dim
          self.num_heads = num_heads
          self.dropout = dropout
          self.head_dim = embed_dim // num_heads
          assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
          self.scaling = self.head_dim ** -0.5
  
          self.in_proj_weight = Parameter(torch.Tensor(3 * embed_dim, embed_dim))
          if bias:
              self.in_proj_bias = Parameter(torch.Tensor(3 * embed_dim))
          else:
              self.register_parameter('in_proj_bias', None)
          self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
  
          if add_bias_kv:
              self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
              self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
          else:
              self.bias_k = self.bias_v = None
  
          self.add_zero_attn = add_zero_attn
  
          self.reset_parameters()
  
          self.onnx_trace = False
  
      def prepare_for_onnx_export_(self):
          self.onnx_trace = True
  
      def reset_parameters(self):
          nn.init.xavier_uniform_(self.in_proj_weight)
          nn.init.xavier_uniform_(self.out_proj.weight)
          if self.in_proj_bias is not None:
              nn.init.constant_(self.in_proj_bias, 0.)
              nn.init.constant_(self.out_proj.bias, 0.)
          if self.bias_k is not None:
              nn.init.xavier_normal_(self.bias_k)
          if self.bias_v is not None:
              nn.init.xavier_normal_(self.bias_v)
  
      def forward(self, query, key, value, key_padding_mask=None, incremental_state=None,
                  need_weights=True, static_kv=False, attn_mask=None):
          """Input shape: Time x Batch x Channel
  
          Self-attention can be implemented by passing in the same arguments for
          query, key and value. Timesteps can be masked by supplying a T x T mask in the
          `attn_mask` argument. Padding elements can be excluded from
          the key by passing a binary ByteTensor (`key_padding_mask`) with shape:
          batch x src_len, where padding elements are indicated by 1s.
          """
  
          qkv_same = query.data_ptr() == key.data_ptr() == value.data_ptr()
          kv_same = key.data_ptr() == value.data_ptr()
  
          tgt_len, bsz, embed_dim = query.size()
          assert embed_dim == self.embed_dim
          assert list(query.size()) == [tgt_len, bsz, embed_dim]
          assert key.size() == value.size()
  
          if incremental_state is not None:
              saved_state = self._get_input_buffer(incremental_state)
              if 'prev_key' in saved_state:
                  # previous time steps are cached - no need to recompute
                  # key and value if they are static
                  if static_kv:
                      assert kv_same and not qkv_same
                      key = value = None
          else:
              saved_state = None
  
          if qkv_same:
              # self-attention
              q, k, v = self.in_proj_qkv(query)
          elif kv_same:
              # encoder-decoder attention
              q = self.in_proj_q(query)
              if key is None:
                  assert value is None
                  k = v = None
              else:
                  k, v = self.in_proj_kv(key)
          else:
              q = self.in_proj_q(query)
              k = self.in_proj_k(key)
              v = self.in_proj_v(value)
  
          q *= self.scaling
  
          if self.bias_k is not None:
              assert self.bias_v is not None
              k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
              v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
              if attn_mask is not None:
                  attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
              if key_padding_mask is not None:
                  key_padding_mask = torch.cat(
                      [key_padding_mask, key_padding_mask.new_zeros(key_padding_mask.size(0), 1)], dim=1)
  
          q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
          if k is not None:
              k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
          if v is not None:
              v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
  
          if saved_state is not None:
              # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
              if 'prev_key' in saved_state:
                  prev_key = saved_state['prev_key'].view(bsz * self.num_heads, -1, self.head_dim)
                  if static_kv:
                      k = prev_key
                  else:
                      k = torch.cat((prev_key, k), dim=1)
              if 'prev_value' in saved_state:
                  prev_value = saved_state['prev_value'].view(bsz * self.num_heads, -1, self.head_dim)
                  if static_kv:
                      v = prev_value
                  else:
                      v = torch.cat((prev_value, v), dim=1)
              saved_state['prev_key'] = k.view(bsz, self.num_heads, -1, self.head_dim)
              saved_state['prev_value'] = v.view(bsz, self.num_heads, -1, self.head_dim)
  
              self._set_input_buffer(incremental_state, saved_state)
  
          src_len = k.size(1)
  
          if key_padding_mask is not None:
              assert key_padding_mask.size(0) == bsz
              assert key_padding_mask.size(1) == src_len
  
          if self.add_zero_attn:
              src_len += 1
              k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
              v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
              if attn_mask is not None:
                  attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)
              if key_padding_mask is not None:
                  key_padding_mask = torch.cat(
                      [key_padding_mask, torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask)], dim=1)
  
          attn_weights = torch.bmm(q, k.transpose(1, 2))
          assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
  
          if attn_mask is not None:
              attn_mask = attn_mask.unsqueeze(0)
              if self.onnx_trace:
                  attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1)
              attn_weights += attn_mask
  
          if key_padding_mask is not None:
              # don't attend to padding symbols
              attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              if self.onnx_trace:
                  attn_weights = torch.where(
                      key_padding_mask.unsqueeze(1).unsqueeze(2),
                      torch.Tensor([float("-Inf")]),
                      attn_weights.float()
                  ).type_as(attn_weights)
              else:
                  attn_weights = attn_weights.float().masked_fill(
                      key_padding_mask.unsqueeze(1).unsqueeze(2),
                      float('-inf'),
                  ).type_as(attn_weights)  # FP16 support: cast to float and back
              attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
              all_inf = torch.isinf(attn_weights).all(dim=-1)
              if all_inf.any():
                  attn_weights = attn_weights.float().masked_fill(
                      all_inf.unsqueeze(-1),
                      0,
                  ).type_as(attn_weights)  # FP16 support: cast to float and back
  
  
          attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(attn_weights)
          attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)
  
          attn = torch.bmm(attn_weights, v)
          assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
          if (self.onnx_trace and attn.size(1) == 1):
              # when ONNX tracing a single decoder step (sequence length == 1)
              # the transpose is a no-op copy before view, thus unnecessary
              attn = attn.contiguous().view(tgt_len, bsz, embed_dim)
          else:
              attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
          attn = self.out_proj(attn)
          if need_weights:
              # average attention weights over heads
              attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              attn_weights = attn_weights.sum(dim=1) / self.num_heads
          else:
              attn_weights = None
  
          return attn, attn_weights
  
      def in_proj_qkv(self, query):
          return self._in_proj(query).chunk(3, dim=-1)
  
      def in_proj_kv(self, key):
          return self._in_proj(key, start=self.embed_dim).chunk(2, dim=-1)
  
      def in_proj_q(self, query):
          return self._in_proj(query, end=self.embed_dim)
  
      def in_proj_k(self, key):
          return self._in_proj(key, start=self.embed_dim, end=2 * self.embed_dim)
  
      def in_proj_v(self, value):
          return self._in_proj(value, start=2 * self.embed_dim)
  
      def _in_proj(self, input, start=0, end=None):
          weight = self.in_proj_weight
          bias = self.in_proj_bias
          weight = weight[start:end, :]
          if bias is not None:
              bias = bias[start:end]
          return F.linear(input, weight, bias)
  
      def reorder_incremental_state(self, incremental_state, new_order):
          """Reorder buffered internal state (for incremental generation)."""
          input_buffer = self._get_input_buffer(incremental_state)
          if input_buffer is not None:
              for k in input_buffer.keys():
                  input_buffer[k] = input_buffer[k].index_select(0, new_order)
              self._set_input_buffer(incremental_state, input_buffer)
  
      def _get_input_buffer(self, incremental_state):
          return utils.get_incremental_state(
              self,
              incremental_state,
              'attn_state',
          ) or {}
  
      def _set_input_buffer(self, incremental_state, buffer):
          utils.set_incremental_state(
              self,
              incremental_state,
              'attn_state',
              buffer,
          )
  
  def ProtectedPositionalEmbedding(
          num_embeddings: int,
          embedding_dim: int,
          padding_idx: int,
          learned: bool = False,
  ):
      if learned:
          # if padding_idx is specified then offset the embedding ids by
          # this index and adjust num_embeddings appropriately
          # TODO: The right place for this offset would be inside
          # LearnedPositionalEmbedding. Move this there for a cleaner implementation.
          if padding_idx is not None:
              num_embeddings = num_embeddings + padding_idx + 1
          m = LearnedPositionalEmbedding(num_embeddings, embedding_dim, padding_idx)
          nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
          if padding_idx is not None:
              nn.init.constant_(m.weight[padding_idx], 0)
      else:
          m = SinusoidalPositionalEmbedding(
              embedding_dim, padding_idx, init_size=num_embeddings + padding_idx + 1,
          )
      return m
  
  
  # class SinusoidalPositionalEmbedding(nn.Module):
  #     """This module produces sinusoidal positional embeddings of any length.
  #
  #     Padding symbols are ignored.
  #     """
  #
  #     def __init__(self, embedding_dim, padding_idx, init_size=1024):
  #         super().__init__()
  #         self.embedding_dim = embedding_dim
  #         self.padding_idx = padding_idx
  #         self.weights = SinusoidalPositionalEmbedding.get_embedding(
  #             init_size,
  #             embedding_dim,
  #         )
  #         self.onnx_trace = False
  #         self.register_buffer('_float_tensor', torch.FloatTensor(1))
  #
  #     def prepare_for_onnx_export_(self):
  #         self.onnx_trace = True
  #
  #     @staticmethod
  #     def get_embedding(num_embeddings, embedding_dim):
  #         """Build sinusoidal embeddings.
  #
  #         This matches the implementation in tensor2tensor, but differs slightly
  #         from the description in Section 3.5 of "Attention Is All You Need".
  #         """
  #         assert embedding_dim % 2 == 0
  #         pos_seq = torch.arange(0, num_embeddings, 1.0, dtype=torch.float)
  #         freq_seq = torch.arange(0, embedding_dim, 2.0, dtype=torch.float)
  #         inv_freq = 1. / (10000 ** (freq_seq / embedding_dim))
  #         sinusoidal_inp = torch.einsum("i, d-> id", pos_seq, inv_freq)
  #         emb = torch.cat([torch.sin(sinusoidal_inp), torch.cos(sinusoidal_inp)], -1)
  #         return emb
  #
  #     def forward(self, input_, incremental_state=None, timestep=None, **kwargs):
  #         """Input is expected to be of size [bsz x seqlen]."""
  #         bsz, seq_len = torch.onnx.operators.shape_as_tensor(input_)
  #         if (seq_len > self.weights.size(0)):
  #             self.weights = SinusoidalPositionalEmbedding.get_embedding(
  #                 seq_len,
  #                 embedding_dim,
  #             )
  #         self.weights = self.weights.to(self._float_tensor)
  #         if incremental_state is not None:
  #             # positions is the same for every token when decoding a single step
  #             pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
  #             if self.onnx_trace:
  #                 return self.weights.index_select(index=pos, dim=0).unsqueeze(1).repeat(bsz, 1, 1)
  #             return self.weights[pos, :].expand(bsz, 1, -1)
  #
  #         weights = self.weights
  #         #get positions
  #         mask = input_.ne(self.padding_idx)
  #         positions = mask.cumsum(-1) - 1
  #         positions = positions * mask.long()
  #         positions = positions.view(-1, )
  #         weights = torch.index_select(weights, 0, positions)
  #         weights = weights.view(bsz, seq_len, -1)
  #         weights = weights * mask.unsqueeze(-1).type_as(weights)
  #         return weights
  #
  #     def max_positions(self):
  #         """Maximum number of supported positions."""
  #         return int(1e5)  # an arbitrary large number
  
  
  class SinusoidalPositionalEmbedding(nn.Module):
      """This module produces sinusoidal positional embeddings of any length.
  
      Padding symbols are ignored.
      """
  
      def __init__(self, embedding_dim, padding_idx, init_size=1024):
          super().__init__()
          self.embedding_dim = embedding_dim
          self.padding_idx = padding_idx
          self.weights = SinusoidalPositionalEmbedding.get_embedding(
              init_size,
              embedding_dim,
          )
          self.onnx_trace = False
          self.register_buffer('_float_tensor', torch.FloatTensor(1))
  
      def prepare_for_onnx_export_(self):
          self.onnx_trace = True
  
      @staticmethod
      def get_embedding(num_embeddings, embedding_dim):
          """Build sinusoidal embeddings.
  
          This matches the implementation in tensor2tensor, but differs slightly
          from the description in Section 3.5 of "Attention Is All You Need".
          """
          half_dim = embedding_dim // 2
          # emb = math.log(10000) / (half_dim - 1)
          emb = math.log(10000) / half_dim
          emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
          emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
          emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
          if embedding_dim % 2 == 1:
              # zero pad
              emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
          return emb
  
      def forward(self, input, incremental_state=None, timestep=None, **kwargs):
          """Input is expected to be of size [bsz x seqlen]."""
          bsz, seq_len = torch.onnx.operators.shape_as_tensor(input)
          max_pos = seq_len
          if self.weights is None or max_pos > self.weights.size(0):
              # recompute/expand embeddings if needed
              self.weights = SinusoidalPositionalEmbedding.get_embedding(
                  max_pos,
                  self.embedding_dim,
              )
          self.weights = self.weights.to(self._float_tensor)
  
          if incremental_state is not None:
              # positions is the same for every token when decoding a single step
              pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
              if self.onnx_trace:
                  return self.weights.index_select(index=-1 + pos, dim=0).unsqueeze(1).repeat(bsz, 1, 1)
              return self.weights[-1 + pos, :].expand(bsz, 1, -1)
  
          positions = make_positions(input, self.padding_idx, onnx_trace=self.onnx_trace)
          if self.onnx_trace:
              flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))
              embedding_shape = torch.cat((bsz.view(1), seq_len.view(1), torch.LongTensor([-1])))
              embeddings = torch.onnx.operators.reshape_from_tensor_shape(flat_embeddings, embedding_shape)
              return embeddings
          return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()
  
      def max_positions(self):
          """Maximum number of supported positions."""
          return int(1e5)  # an arbitrary large number
  
  
  def make_positions(tensor, padding_idx, onnx_trace=False):
      """Replace non-padding symbols with their position numbers.
  
      Position numbers begin at padding_idx+1. Padding symbols are ignored.
      """
      # The series of casts and type-conversions here are carefully
      # balanced to both work with ONNX export and XLA. In particular XLA
      # prefers ints, cumsum defaults to output longs, and ONNX doesn't know
      # how to handle the dtype kwarg in cumsum.
      mask = tensor.ne(padding_idx).int()
      return (
          (torch.cumsum(mask, dim=1) - 1).type_as(mask) * mask
      ).long()
  
  
  
apis: torch.nn.init.xavier_uniform_ torch.cumsum torch.bmm torch.cat torch.onnx.operators.reshape_from_tensor_shape torch.nn.functional.dropout torch.nn.init.xavier_normal_ torch.cos torch.onnx.operators.shape_as_tensor torch.nn.init.normal_ torch.isinf torch.sin torch.arange torch.Tensor torch.FloatTensor torch.nn.Linear torch.nn.init.constant_ torch.nn.functional.linear torch.zeros torch.LongTensor