text: 
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import math
  import torch.onnx.operators
  
  from typing import Callable, Optional, Dict, Tuple, Any, NamedTuple, List
  from torch import Tensor, nn
  from torch.nn import Parameter
  import uuid
  
  try:
      from apex.normalization import FusedLayerNorm as _FusedLayerNorm
  
      has_fused_layernorm = True
  
      class FusedLayerNorm(_FusedLayerNorm):
          @torch.jit.unused
          def forward(self, x):
              if not x.is_cuda:
                  return super().forward(x)
              else:
                  with torch.cuda.device(x.device):
                      return super().forward(x)
  
  except ImportError:
      has_fused_layernorm = False
  from transformers.utils import logging
  
  def LayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True, export=False):
      if not export and torch.cuda.is_available() and has_fused_layernorm:
          return FusedLayerNorm(normalized_shape, eps, elementwise_affine)
      return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)
  
  
  # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  
  
  def softmax(x, dim: int, onnx_trace: bool = False):
      if onnx_trace:
          return F.softmax(x.float(), dim=dim)
      else:
          return F.softmax(x, dim=dim, dtype=torch.float32)
  
  def log_softmax(x, dim: int, onnx_trace: bool = False):
      if onnx_trace:
          return F.log_softmax(x.float(), dim=dim)
      else:
          return F.log_softmax(x, dim=dim, dtype=torch.float32)
  
  def gelu_accurate(x):
      if not hasattr(gelu_accurate, "_a"):
          gelu_accurate._a = math.sqrt(2 / math.pi)
      return (
          0.5 * x * (1 + torch.tanh(gelu_accurate._a * (x + 0.044715 * torch.pow(x, 3))))
      )
  
  
  def gelu(x: torch.Tensor) -> torch.Tensor:
      return torch.nn.functional.gelu(x.float()).type_as(x)
  
  
  def get_activation_fn(activation: str) -> Callable:
      """ Returns the activation function corresponding to `activation` """
      if activation == "relu":
          return F.relu
      elif activation == "gelu":
          return gelu
      elif activation == "gelu_fast":
          deprecation_warning(
              "--activation-fn=gelu_fast has been renamed to gelu_accurate"
          )
          return gelu_accurate
      elif activation == "gelu_accurate":
          return gelu_accurate
      elif activation == "tanh":
          return torch.tanh
      elif activation == "linear":
          return lambda x: x
      else:
          raise RuntimeError("--activation-fn {} not supported".format(activation))
  
  
  def make_positions(tensor, padding_idx: int, onnx_trace: bool = 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).type_as(mask) * mask).long() + padding_idx
  
  
  
  class LayerDropModuleList(nn.ModuleList):
      """
      A LayerDrop implementation based on :class:`torch.nn.ModuleList`.
  
      We refresh the choice of which layers to drop every time we iterate
      over the LayerDropModuleList instance. During evaluation we always
      iterate over all layers.
  
      Usage::
  
          layers = LayerDropList(p=0.5, modules=[layer1, layer2, layer3])
          for layer in layers:  # this might iterate over layers 1 and 3
              x = layer(x)
          for layer in layers:  # this might iterate over all layers
              x = layer(x)
          for layer in layers:  # this might not iterate over any layers
              x = layer(x)
  
      Args:
          p (float): probability of dropping out each layer
          modules (iterable, optional): an iterable of modules to add
      """
  
      def __init__(self, p, modules=None):
          super().__init__(modules)
          self.p = p
  
      def __iter__(self):
          dropout_probs = torch.empty(len(self)).uniform_()
          for i, m in enumerate(super().__iter__()):
              if not self.training or (dropout_probs[i] > self.p):
                  yield m
  
  
  class LearnedPositionalEmbedding(nn.Embedding):
      """
      This module learns positional embeddings up to a fixed maximum size.
      Padding ids are ignored by either offsetting based on padding_idx
      or by setting padding_idx to None and ensuring that the appropriate
      position ids are passed to the forward function.
      """
  
      def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int):
          super().__init__(num_embeddings, embedding_dim, padding_idx)
          self.onnx_trace = False
          if self.padding_idx is not None:
              self.max_positions = self.num_embeddings - self.padding_idx - 1
          else:
              self.max_positions = self.num_embeddings
  
      def forward(
          self,
          input: Tensor,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          positions: Optional[Tensor] = None,
      ):
          """Input is expected to be of size [bsz x seqlen]."""
          assert (positions is None) or (
              self.padding_idx is None
          ), "If positions is pre-computed then padding_idx should not be set."
  
          if positions is None:
              if incremental_state is not None:
                  # positions is the same for every token when decoding a single step
                  # Without the int() cast, it doesn't work in some cases when exporting to ONNX
                  positions = torch.zeros(
                      (1, 1), device=input.device, dtype=input.dtype
                  ).fill_(int(self.padding_idx + input.size(1)))
              else:
                  positions = make_positions(
                      input, self.padding_idx, onnx_trace=self.onnx_trace
                  )
          return F.embedding(
              positions,
              self.weight,
              self.padding_idx,
              self.max_norm,
              self.norm_type,
              self.scale_grad_by_freq,
              self.sparse,
          )
  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, padding_idx
          )
          self.onnx_trace = False
          self.register_buffer("_float_tensor", torch.FloatTensor(1))
          self.max_positions = int(1e5)
  
      def prepare_for_onnx_export_(self):
          self.onnx_trace = True
  
      @staticmethod
      def get_embedding(
          num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None
      ):
          """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 = 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)
          if padding_idx is not None:
              emb[padding_idx, :] = 0
          return emb
  
      def forward(
          self,
          input,
          incremental_state: Optional[Any] = None,
          timestep: Optional[Tensor] = None,
          positions: Optional[Any] = None,
      ):
          """Input is expected to be of size [bsz x seqlen]."""
          bspair = torch.onnx.operators.shape_as_tensor(input)
          bsz, seq_len = bspair[0], bspair[1]
          max_pos = self.padding_idx + 1 + 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.padding_idx
              )
          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=self.padding_idx + pos, dim=0)
                      .unsqueeze(1)
                      .repeat(bsz, 1, 1)
                  )
              return self.weights[self.padding_idx + 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.tensor([-1], dtype=torch.long))
              )
              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 PositionalEmbedding(
          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
  
  
  
  
  
  
  logger = logging.get_logger(__name__)
  
  
  class FairseqDropout(nn.Module):
  
      def __init__(self, p, module_name=None):
          super().__init__()
          self.p = p
          self.module_name = module_name
          self.apply_during_inference = False
  
      def forward(self, x, inplace: bool = False):
          if self.training or self.apply_during_inference:
              return F.dropout(x, p=self.p, training=True, inplace=inplace)
          else:
              return x
  
      def make_generation_fast_(
          self,
          name: str,
          retain_dropout: bool = False,
          retain_dropout_modules: Optional[List[str]] = None,
          **kwargs
      ):
          if retain_dropout:
              if retain_dropout_modules is not None and self.module_name is None:
                  logger.warning(
                      'Cannot enable dropout during inference for module {} '
                      'because module_name was not set'.format(name)
                  )
              elif (
                  retain_dropout_modules is None  # if None, apply to all modules
                  or self.module_name in retain_dropout_modules
              ):
                  logger.info(
                      'Enabling dropout during inference for module: {}'.format(name)
                  )
                  self.apply_during_inference = True
              else:
                  logger.info('Disabling dropout for module: {}'.format(name))
  
  
  
  
  
  class FairseqIncrementalState(object):
  
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self.init_incremental_state()
  
      def init_incremental_state(self):
          self._incremental_state_id = str(uuid.uuid4())
  
      def _get_full_incremental_state_key(self, key: str) -> str:
          return "{}.{}".format(self._incremental_state_id, key)
  
      def get_incremental_state(
          self,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
          key: str,
      ) -> Optional[Dict[str, Optional[Tensor]]]:
          """Helper for getting incremental state for an nn.Module."""
          full_key = self._get_full_incremental_state_key(key)
          if incremental_state is None or full_key not in incremental_state:
              return None
          return incremental_state[full_key]
  
      def set_incremental_state(
          self,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
          key: str,
          value: Dict[str, Optional[Tensor]],
      ) -> Optional[Dict[str, Dict[str, Optional[Tensor]]]]:
          """Helper for setting incremental state for an nn.Module."""
          if incremental_state is not None:
              full_key = self._get_full_incremental_state_key(key)
              incremental_state[full_key] = value
          return incremental_state
  
  
  def with_incremental_state(cls):
      cls.__bases__ = (FairseqIncrementalState,) + tuple(b for b in cls.__bases__ if b != FairseqIncrementalState)
      return cls
  
  
  
  
  class SelfMultiheadAttention(nn.Module):
      """Multi-headed attention.
      See "Attention Is All You Need" for more details.
      """
  
      def __init__(
          self,
          embed_dim,
          num_heads,
          kdim=None,
          vdim=None,
          dropout=0.0,
          bias=True,
          add_bias_kv=False,
          add_zero_attn=False,
          self_attention=False,
          encoder_decoder_attention=False,
          q_noise=0.0,
          qn_block_size=8,
          scaling_factor=1,
      ):
          super().__init__()
          self.embed_dim = embed_dim
          self.kdim = kdim if kdim is not None else embed_dim
          self.vdim = vdim if vdim is not None else embed_dim
          self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
  
          self.num_heads = num_heads
          self.dropout_module = FairseqDropout(
              dropout, module_name=self.__class__.__name__
          )
  
          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.self_attention = self_attention
          self.encoder_decoder_attention = encoder_decoder_attention
  
          assert not self.self_attention or self.qkv_same_dim, (
              "Self-attention requires query, key and " "value to be of the same size"
          )
  
          self.k_proj = quant_noise(nn.Linear(self.kdim, embed_dim, bias=bias), q_noise, qn_block_size)
          self.v_proj = quant_noise(nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size)
          self.q_proj = quant_noise(nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size)
  
          self.out_proj = quant_noise(nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size)
  
          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
          self.tpu = False
  
      def reset_parameters(self):
          if self.qkv_same_dim:
              # Empirically observed the convergence to be much better with
              # the scaled initialization
              nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
              nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
              nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
          else:
              nn.init.xavier_uniform_(self.k_proj.weight)
              nn.init.xavier_uniform_(self.v_proj.weight)
              nn.init.xavier_uniform_(self.q_proj.weight)
  
          nn.init.xavier_uniform_(self.out_proj.weight)
          if self.out_proj.bias is not None:
              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 prepare_for_onnx_export_(self):
          self.onnx_trace = True
  
      def prepare_for_tpu_(self, **kwargs):
          self.tpu = True
  
      def forward(
          self,
          query,
          key_padding_mask: Optional[Tensor] = None,
          need_weights: bool = True,
          attn_mask: Optional[Tensor] = None,
          attn_bias: Optional[Tensor] = None,
          before_softmax: bool = False,
          need_head_weights: bool = False,
      ) -> Tuple[Tensor, Optional[Tensor]]:
          """Input shape: Time x Batch x Channel
          Args:
              key_padding_mask (ByteTensor, optional): mask to exclude
                  keys that are pads, of shape `(batch, src_len)`, where
                  padding elements are indicated by 1s.
              need_weights (bool, optional): return the attention weights,
                  averaged over heads (default: False).
              attn_mask (ByteTensor, optional): typically used to
                  implement causal attention, where the mask prevents the
                  attention from looking forward in time (default: None).
              before_softmax (bool, optional): return the raw attention
                  weights and values before the attention softmax.
              need_head_weights (bool, optional): return the attention
                  weights for each head. Implies *need_weights*. Default:
                  return the average attention weights over all heads.
          """
          if need_head_weights:
              need_weights = True
  
          tgt_len, bsz, embed_dim = query.size()
          assert embed_dim == self.embed_dim
          assert list(query.size()) == [tgt_len, bsz, embed_dim]
  
          q = self.q_proj(query)
          k = self.k_proj(query)
          v = self.v_proj(query)
          q *= self.scaling
  
          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)
              )
  
          assert k is not None
          src_len = k.size(1)
  
          # This is part of a workaround to get around fork/join parallelism
          # not supporting Optional types.
          if key_padding_mask is not None and key_padding_mask.dim() == 0:
              key_padding_mask = None
  
          if key_padding_mask is not None:
              assert key_padding_mask.size(0) == bsz
              assert key_padding_mask.size(1) == src_len
  
  
          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 not self.tpu:
                  # attn_weights = attn_weights.masked_fill(
                  #     key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
                  #     float("-inf")
                  # )
                  attn_weights = attn_weights.masked_fill(
                      key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
                      float(-1e4)
                  )
              else:
                  attn_weights = attn_weights.transpose(0, 2)
                  attn_weights = attn_weights.masked_fill(key_padding_mask, float('-inf'))
                  attn_weights = attn_weights.transpose(0, 2)
              attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
  
          if attn_bias is not None:
              attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              attn_weights += attn_bias
              attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
  
          if before_softmax:
              return attn_weights, v
  
          #print('????attn_weights: ',attn_weights)
          attn_weights_float = softmax(
              attn_weights, dim=-1, onnx_trace=self.onnx_trace
          )
          attn_weights = attn_weights_float.type_as(attn_weights)
          attn_probs = self.dropout_module(attn_weights)
          # print('????attn_probs: ',attn_probs)
          # assert 1==0
  
          assert v is not None
          attn = torch.bmm(attn_probs, 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)
          attn_weights: Optional[Tensor] = None
          if need_weights:
              attn_weights = attn_weights_float.view(
                  bsz, self.num_heads, tgt_len, src_len
              ).transpose(1, 0)
              if not need_head_weights:
                  # average attention weights over heads
                  attn_weights = attn_weights.mean(dim=0)
  
          return attn, attn_weights
  
  
  @with_incremental_state
  class MultiheadAttention(nn.Module):
      """Multi-headed attention.
  
      See "Attention Is All You Need" for more details.
      """
  
      def __init__(
          self,
          embed_dim,
          num_heads,
          kdim=None,
          vdim=None,
          dropout=0.0,
          bias=True,
          add_bias_kv=False,
          add_zero_attn=False,
          self_attention=False,
          encoder_decoder_attention=False,
          q_noise=0.0,
          qn_block_size=8,
      ):
          super().__init__()
          self.embed_dim = embed_dim
          self.kdim = kdim if kdim is not None else embed_dim
          self.vdim = vdim if vdim is not None else embed_dim
          self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
  
          self.num_heads = num_heads
          self.dropout_module = FairseqDropout(
              dropout, module_name=self.__class__.__name__
          )
  
          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.self_attention = self_attention
          self.encoder_decoder_attention = encoder_decoder_attention
  
          assert not self.self_attention or self.qkv_same_dim, (
              "Self-attention requires query, key and " "value to be of the same size"
          )
  
          self.k_proj = quant_noise(nn.Linear(self.kdim, embed_dim, bias=bias), q_noise, qn_block_size)
          self.v_proj = quant_noise(nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size)
          self.q_proj = quant_noise(nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size)
  
          self.out_proj = quant_noise(nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size)
  
          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
          self.tpu = False
  
      def prepare_for_onnx_export_(self):
          self.onnx_trace = True
  
      def prepare_for_tpu_(self, **kwargs):
          self.tpu = True
  
      def reset_parameters(self):
          if self.qkv_same_dim:
              # Empirically observed the convergence to be much better with
              # the scaled initialization
              nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
              nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
              nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
          else:
              nn.init.xavier_uniform_(self.k_proj.weight)
              nn.init.xavier_uniform_(self.v_proj.weight)
              nn.init.xavier_uniform_(self.q_proj.weight)
  
          nn.init.xavier_uniform_(self.out_proj.weight)
          if self.out_proj.bias is not None:
              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: Optional[Tensor],
          value: Optional[Tensor],
          key_padding_mask: Optional[Tensor] = None,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          need_weights: bool = True,
          static_kv: bool = False,
          attn_mask: Optional[Tensor] = None,
          before_softmax: bool = False,
          need_head_weights: bool = False,
      ) -> Tuple[Tensor, Optional[Tensor]]:
          """Input shape: Time x Batch x Channel
  
          Args:
              key_padding_mask (ByteTensor, optional): mask to exclude
                  keys that are pads, of shape `(batch, src_len)`, where
                  padding elements are indicated by 1s.
              need_weights (bool, optional): return the attention weights,
                  averaged over heads (default: False).
              attn_mask (ByteTensor, optional): typically used to
                  implement causal attention, where the mask prevents the
                  attention from looking forward in time (default: None).
              before_softmax (bool, optional): return the raw attention
                  weights and values before the attention softmax.
              need_head_weights (bool, optional): return the attention
                  weights for each head. Implies *need_weights*. Default:
                  return the average attention weights over all heads.
          """
          if need_head_weights:
              need_weights = True
  
          tgt_len, bsz, embed_dim = query.size()
          assert embed_dim == self.embed_dim
          assert list(query.size()) == [tgt_len, bsz, embed_dim]
  
          if (
              not self.onnx_trace
              and not self.tpu  # don't use PyTorch version on TPUs
              and incremental_state is None
              and not static_kv
              # A workaround for quantization to work. Otherwise JIT compilation
              # treats bias in linear module as method.
              and not torch.jit.is_scripting()
          ):
              assert key is not None and value is not None
              # print('okok')
              # assert 1==0
              return F.multi_head_attention_forward(
                  query,
                  key,
                  value,
                  self.embed_dim,
                  self.num_heads,
                  torch.empty([0]),
                  torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)),
                  self.bias_k,
                  self.bias_v,
                  self.add_zero_attn,
                  self.dropout_module.p,
                  self.out_proj.weight,
                  self.out_proj.bias,
                  self.training or self.dropout_module.apply_during_inference,
                  key_padding_mask,
                  need_weights,
                  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,
              )
          #print('???',self.onnx_trace, self.tpu,incremental_state,static_kv,torch.jit.is_scripting())
          if incremental_state is not None:
              saved_state = self._get_input_buffer(incremental_state)
              if saved_state is not None and "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 self.encoder_decoder_attention and not self.self_attention
                      key = value = None
          else:
              saved_state = None
  
          if self.self_attention:
              q = self.q_proj(query)
              k = self.k_proj(query)
              v = self.v_proj(query)
          elif self.encoder_decoder_attention:
              # encoder-decoder attention
              q = self.q_proj(query)
              if key is None:
                  assert value is None
                  k = v = None
              else:
                  k = self.k_proj(key)
                  v = self.v_proj(key)
  
          else:
              assert key is not None and value is not None
              q = self.q_proj(query)
              k = self.k_proj(key)
              v = self.v_proj(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)
              )
          #print('???qkv',q.shape,k.shape,v.shape,bsz,self.num_heads, self.head_dim)
          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"]
                  assert _prev_key is not None
                  prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
                  if static_kv:
                      k = prev_key
                  else:
                      assert k is not None
                      k = torch.cat([prev_key, k], dim=1)
              if "prev_value" in saved_state:
                  _prev_value = saved_state["prev_value"]
                  assert _prev_value is not None
                  prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
                  if static_kv:
                      v = prev_value
                  else:
                      assert v is not None
                      v = torch.cat([prev_value, v], dim=1)
              prev_key_padding_mask: Optional[Tensor] = None
              if "prev_key_padding_mask" in saved_state:
                  prev_key_padding_mask = saved_state["prev_key_padding_mask"]
              assert k is not None and v is not None
              key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
                  key_padding_mask=key_padding_mask,
                  prev_key_padding_mask=prev_key_padding_mask,
                  batch_size=bsz,
                  src_len=k.size(1),
                  static_kv=static_kv,
              )
  
              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)
              saved_state["prev_key_padding_mask"] = key_padding_mask
              # In this branch incremental_state is never None
              assert incremental_state is not None
              incremental_state = self._set_input_buffer(incremental_state, saved_state)
          assert k is not None
          src_len = k.size(1)
  
          # This is part of a workaround to get around fork/join parallelism
          # not supporting Optional types.
          if key_padding_mask is not None and key_padding_mask.dim() == 0:
              key_padding_mask = None
  
          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:
              assert v is not None
              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))
          attn_weights = MultiheadAttention.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
  
          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 not self.tpu:
                  attn_weights = attn_weights.masked_fill(
                      key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
                      float("-inf")
                  )
              else:
                  attn_weights = attn_weights.transpose(0, 2)
                  attn_weights = attn_weights.masked_fill(key_padding_mask, float('-inf'))
                  attn_weights = attn_weights.transpose(0, 2)
              attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
  
          if before_softmax:
              return attn_weights, v
  
          attn_weights_float = softmax(
              attn_weights, dim=-1, onnx_trace=self.onnx_trace
          )
          attn_weights = attn_weights_float.type_as(attn_weights)
          attn_probs = self.dropout_module(attn_weights)
  
          assert v is not None
          attn = torch.bmm(attn_probs, 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)
          attn_weights: Optional[Tensor] = None
          if need_weights:
              attn_weights = attn_weights_float.view(
                  bsz, self.num_heads, tgt_len, src_len
              ).transpose(1, 0)
              if not need_head_weights:
                  # average attention weights over heads
                  attn_weights = attn_weights.mean(dim=0)
  
          return attn, attn_weights
  
      @staticmethod
      def _append_prev_key_padding_mask(
          key_padding_mask: Optional[Tensor],
          prev_key_padding_mask: Optional[Tensor],
          batch_size: int,
          src_len: int,
          static_kv: bool,
      ) -> Optional[Tensor]:
          # saved key padding masks have shape (bsz, seq_len)
          if prev_key_padding_mask is not None and static_kv:
              new_key_padding_mask = prev_key_padding_mask
          elif prev_key_padding_mask is not None and key_padding_mask is not None:
              new_key_padding_mask = torch.cat(
                  [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
              )
          # During incremental decoding, as the padding token enters and
          # leaves the frame, there will be a time when prev or current
          # is None
          elif prev_key_padding_mask is not None:
              filler = torch.zeros(
                  (batch_size, src_len - prev_key_padding_mask.size(1)),
                  device=prev_key_padding_mask.device,
              )
              new_key_padding_mask = torch.cat(
                  [prev_key_padding_mask.float(), filler.float()], dim=1
              )
          elif key_padding_mask is not None:
              filler = torch.zeros(
                  (batch_size, src_len - key_padding_mask.size(1)),
                  device=key_padding_mask.device,
              )
              new_key_padding_mask = torch.cat(
                  [filler.float(), key_padding_mask.float()], dim=1
              )
          else:
              new_key_padding_mask = prev_key_padding_mask
          return new_key_padding_mask
  
      @torch.jit.export
      def reorder_incremental_state(
          self, incremental_state: Dict[str, Dict[str, Optional[Tensor]]], new_order: Tensor
      ):
          """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]
                  if input_buffer_k is not None:
                      if self.encoder_decoder_attention and input_buffer_k.size(0) == new_order.size(0):
                          break
                      input_buffer[k] = input_buffer_k.index_select(0, new_order)
              incremental_state = self._set_input_buffer(incremental_state, input_buffer)
          return incremental_state
  
      def _get_input_buffer(
          self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]
      ) -> Dict[str, Optional[Tensor]]:
          result = self.get_incremental_state(incremental_state, "attn_state")
          if result is not None:
              return result
          else:
              empty_result: Dict[str, Optional[Tensor]] = {}
              return empty_result
  
      def _set_input_buffer(
          self,
          incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
          buffer: Dict[str, Optional[Tensor]],
      ):
          return self.set_incremental_state(incremental_state, "attn_state", buffer)
  
      def apply_sparse_mask(attn_weights, tgt_len: int, src_len: int, bsz: int):
          return attn_weights
  
      def upgrade_state_dict_named(self, state_dict, name):
          prefix = name + "." if name != "" else ""
          items_to_add = {}
          keys_to_remove = []
          for k in state_dict.keys():
              if k.endswith(prefix + "in_proj_weight"):
                  # in_proj_weight used to be q + k + v with same dimensions
                  dim = int(state_dict[k].shape[0] / 3)
                  items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim]
                  items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim]
                  items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :]
  
                  keys_to_remove.append(k)
  
                  k_bias = prefix + "in_proj_bias"
                  if k_bias in state_dict.keys():
                      dim = int(state_dict[k].shape[0] / 3)
                      items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim]
                      items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][
                          dim : 2 * dim
                      ]
                      items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :]
  
                      keys_to_remove.append(prefix + "in_proj_bias")
  
          for k in keys_to_remove:
              del state_dict[k]
  
          for key, value in items_to_add.items():
              state_dict[key] = value
  
  
  class AdaptiveSoftmax(nn.Module):
      """
      This is an implementation of the efficient softmax approximation for
      graphical processing units (GPU), described in the paper "Efficient softmax
      approximation for GPUs" (http://arxiv.org/abs/1609.04309).
      """
  
      def __init__(self, vocab_size, input_dim, cutoff, dropout, factor=4., adaptive_inputs=None, tie_proj=False,
                   q_noise=0, qn_block_size=8):
          super().__init__()
  
          if vocab_size > cutoff[-1]:
              cutoff = cutoff + [vocab_size]
          else:
              assert vocab_size == cutoff[
                  -1], 'cannot specify cutoff larger than vocab size'
  
          output_dim = cutoff[0] + len(cutoff) - 1
  
          self.vocab_size = vocab_size
          self.cutoff = cutoff
          self.dropout_module = FairseqDropout(dropout, module_name=self.__class__.__name__)
          self.input_dim = input_dim
          self.factor = factor
          self.q_noise = q_noise
          self.qn_block_size = qn_block_size
  
          self.lsm = nn.LogSoftmax(dim=1)
  
          if adaptive_inputs is not None:
              self.head = TiedHeadModule(adaptive_inputs.weights_for_band(0), input_dim, len(cutoff) - 1, self.q_noise, self.qn_block_size)
          else:
              self.head = quant_noise(nn.Linear(input_dim, output_dim, bias=False), self.q_noise, self.qn_block_size)
  
          self._make_tail(adaptive_inputs, tie_proj)
  
          def init_weights(m):
              if hasattr(m, 'weight') and not isinstance(m, TiedLinear) and not isinstance(m, TiedHeadModule):
                  nn.init.xavier_uniform_(m.weight)
  
          self.apply(init_weights)
  
          self.register_buffer('version', torch.LongTensor([1]))
  
      def _make_tail(self, adaptive_inputs=None, tie_proj=False):
          self.tail = nn.ModuleList()
          for i in range(len(self.cutoff) - 1):
              dim = int(self.input_dim // self.factor ** (i + 1))
  
              tied_emb, tied_proj = adaptive_inputs.weights_for_band(i + 1) \
                  if adaptive_inputs is not None else (None, None)
  
              if tied_proj is not None:
                  if tie_proj:
                      proj = quant_noise(TiedLinear(tied_proj, transpose=True), self.q_noise, self.qn_block_size)
                  else:
                      proj = quant_noise(nn.Linear(tied_proj.size(0), tied_proj.size(1), bias=False), self.q_noise, self.qn_block_size)
              else:
                  proj = quant_noise(nn.Linear(self.input_dim, dim, bias=False), self.q_noise, self.qn_block_size)
  
              if tied_emb is None:
                  out_proj = nn.Linear(dim, self.cutoff[i + 1] - self.cutoff[i], bias=False)
              else:
                  out_proj = TiedLinear(tied_emb, transpose=False)
  
              m = nn.Sequential(
                  proj,
                  nn.Dropout(self.dropout_module.p),
                  quant_noise(out_proj, self.q_noise, self.qn_block_size),
              )
  
              self.tail.append(m)
  
      def upgrade_state_dict_named(self, state_dict, name):
          version_name = name + '.version'
          if version_name not in state_dict:
              raise Exception('This version of the model is no longer supported')
  
      def adapt_target(self, target):
          """
          In order to be efficient, the AdaptiveSoftMax does not compute the
          scores for all the word of the vocabulary for all the examples. It is
          thus necessary to call the method adapt_target of the AdaptiveSoftMax
          layer inside each forward pass.
          """
  
          target = target.view(-1)
          new_target = [target.clone()]
          target_idxs = []
  
          for i in range(len(self.cutoff) - 1):
              mask = target.ge(self.cutoff[i]).mul(target.lt(self.cutoff[i + 1]))
              new_target[0][mask] = self.cutoff[0] + i
  
              if mask.any():
                  target_idxs.append(mask.nonzero().squeeze(1))
                  new_target.append(target[mask].add(-self.cutoff[i]))
              else:
                  target_idxs.append(None)
                  new_target.append(None)
  
          return new_target, target_idxs
  
      def forward(self, input, target):
          """
          Args:
              input: (b x t x d)
              target: (b x t)
          Returns:
              2 lists: output for each cutoff section and new targets by cut off
          """
  
          input = input.contiguous().view(-1, input.size(-1))
          input = self.dropout_module(input)
  
          new_target, target_idxs = self.adapt_target(target)
          output = [self.head(input)]
  
          for i in range(len(target_idxs)):
              if target_idxs[i] is not None:
                  output.append(self.tail[i](input.index_select(0, target_idxs[i])))
              else:
                  output.append(None)
  
          return output, new_target
  
      def get_log_prob(self, input, target):
          """
          Computes the log probabilities for all the words of the vocabulary,
          given a 2D tensor of hidden vectors.
          """
  
          bsz, length, dim = input.size()
          input = input.contiguous().view(-1, dim)
  
          if target is not None:
              _, target_idxs = self.adapt_target(target)
          else:
              target_idxs = None
  
          head_y = self.head(input)
          log_probs = head_y.new_zeros(input.size(0), self.vocab_size)
  
          head_sz = self.cutoff[0] + len(self.tail)
          log_probs[:, :head_sz] = self.lsm(head_y)
          tail_priors = log_probs[:, self.cutoff[0]: head_sz].clone()
  
          for i in range(len(self.tail)):
              start = self.cutoff[i]
              end = self.cutoff[i + 1]
  
              if target_idxs is None:
                  tail_out = log_probs[:, start:end]
                  tail_out.copy_(self.tail[i](input))
                  log_probs[:, start:end] = self.lsm(tail_out).add_(tail_priors[:, i, None])
              elif target_idxs[i] is not None:
                  idxs = target_idxs[i]
                  tail_out = log_probs[idxs, start:end]
                  tail_out.copy_(self.tail[i](input[idxs]))
                  log_probs[idxs, start:end] = self.lsm(tail_out).add_(tail_priors[idxs, i, None])
  
          log_probs = log_probs.view(bsz, length, -1)
          return log_probs
  
  
  
  
  class TransformerDecoderLayer(nn.Module):
      """Decoder layer block.
  
      In the original paper each operation (multi-head attention, encoder
      attention or FFN) is postprocessed with: `dropout -> add residual ->
      layernorm`. In the tensor2tensor code they suggest that learning is more
      robust when preprocessing each layer with layernorm and postprocessing with:
      `dropout -> add residual`. We default to the approach in the paper, but the
      tensor2tensor approach can be enabled by setting
      *args.decoder_normalize_before* to ``True``.
  
      Args:
          args (argparse.Namespace): parsed command-line arguments
          no_encoder_attn (bool, optional): whether to attend to encoder outputs
              (default: False).
      """
  
      def __init__(
          self, args, no_encoder_attn=False, add_bias_kv=False, add_zero_attn=False
      ):
          super().__init__()
          self.embed_dim = args.decoder_embed_dim
          self.dropout_module = FairseqDropout(args.dropout, module_name=self.__class__.__name__)
          self.quant_noise = getattr(args, "quant_noise_pq", 0)
          self.quant_noise_block_size = getattr(args, "quant_noise_pq_block_size", 8)
  
          self.cross_self_attention = getattr(args, "cross_self_attention", False)
  
          self.self_attn = self.build_self_attention(
              self.embed_dim,
              args,
              add_bias_kv=add_bias_kv,
              add_zero_attn=add_zero_attn,
          )
          self.activation_fn = get_activation_fn(
              activation=getattr(args, "activation_fn", "relu")
          )
          activation_dropout_p = getattr(args, "activation_dropout", 0)
          if activation_dropout_p == 0:
              # for backwards compatibility with models that use args.relu_dropout
              activation_dropout_p = getattr(args, "relu_dropout", 0)
          self.activation_dropout_module = FairseqDropout(
              float(activation_dropout_p), module_name=self.__class__.__name__)
          self.normalize_before = args.decoder_normalize_before
  
          # use layerNorm rather than FusedLayerNorm for exporting.
          # char_inputs can be used to determint this.
          # TODO  remove this once we update apex with the fix
          export = getattr(args, "char_inputs", False)
          self.self_attn_layer_norm = LayerNorm(self.embed_dim, export=export)
  
          if no_encoder_attn:
              self.encoder_attn = None
              self.encoder_attn_layer_norm = None
          else:
              self.encoder_attn = self.build_encoder_attention(self.embed_dim, args)
              self.encoder_attn_layer_norm = LayerNorm(self.embed_dim, export=export)
  
          self.fc1 = self.build_fc1(
              self.embed_dim, args.decoder_ffn_embed_dim, self.quant_noise, self.quant_noise_block_size
          )
          self.fc2 = self.build_fc2(
              args.decoder_ffn_embed_dim, self.embed_dim, self.quant_noise, self.quant_noise_block_size
          )
  
          self.final_layer_norm = LayerNorm(self.embed_dim, export=export)
          self.need_attn = True
  
          self.onnx_trace = False
  
      def build_fc1(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(nn.Linear(input_dim, output_dim), q_noise, qn_block_size)
  
      def build_fc2(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(nn.Linear(input_dim, output_dim), q_noise, qn_block_size)
  
      def build_self_attention(self, embed_dim, args, add_bias_kv=False, add_zero_attn=False):
          return MultiheadAttention(
              embed_dim,
              args.decoder_attention_heads,
              dropout=args.attention_dropout,
              add_bias_kv=add_bias_kv,
              add_zero_attn=add_zero_attn,
              self_attention=not getattr(args, "cross_self_attention", False),
              q_noise=self.quant_noise,
              qn_block_size=self.quant_noise_block_size,
          )
  
      def build_encoder_attention(self, embed_dim, args):
          return MultiheadAttention(
              embed_dim,
              args.decoder_attention_heads,
              kdim=getattr(args, "encoder_embed_dim", None),
              vdim=getattr(args, "encoder_embed_dim", None),
              dropout=args.attention_dropout,
              encoder_decoder_attention=True,
              q_noise=self.quant_noise,
              qn_block_size=self.quant_noise_block_size,
          )
  
      def prepare_for_onnx_export_(self):
          self.onnx_trace = True
  
      def forward(
          self,
          x,
          encoder_out: Optional[torch.Tensor] = None,
          encoder_padding_mask: Optional[torch.Tensor] = None,
          incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
          prev_self_attn_state: Optional[List[torch.Tensor]] = None,
          prev_attn_state: Optional[List[torch.Tensor]] = None,
          self_attn_mask: Optional[torch.Tensor] = None,
          self_attn_padding_mask: Optional[torch.Tensor] = None,
          need_attn: bool = False,
          need_head_weights: bool = False,
          keep_residual: int =0,
      ):
          """
          Args:
              x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
              encoder_padding_mask (ByteTensor, optional): binary
                  ByteTensor of shape `(batch, src_len)` where padding
                  elements are indicated by ``1``.
              need_attn (bool, optional): return attention weights
              need_head_weights (bool, optional): return attention weights
                  for each head (default: return average over heads).
  
          Returns:
              encoded output of shape `(seq_len, batch, embed_dim)`
          """
          if need_head_weights:
              need_attn = True
  
          residual = x
          if self.normalize_before:
              x = self.self_attn_layer_norm(x)
          if prev_self_attn_state is not None:
              prev_key, prev_value = prev_self_attn_state[:2]
              saved_state: Dict[str, Optional[Tensor]] = {
                  "prev_key": prev_key,
                  "prev_value": prev_value,
              }
              if len(prev_self_attn_state) >= 3:
                  saved_state["prev_key_padding_mask"] = prev_self_attn_state[2]
              assert incremental_state is not None
              self.self_attn._set_input_buffer(incremental_state, saved_state)
          _self_attn_input_buffer = self.self_attn._get_input_buffer(incremental_state)
          if self.cross_self_attention and not (
              incremental_state is not None
              and _self_attn_input_buffer is not None
              and "prev_key" in _self_attn_input_buffer
          ):
              if self_attn_mask is not None:
                  assert encoder_out is not None
                  self_attn_mask = torch.cat(
                      (x.new_zeros(x.size(0), encoder_out.size(0)), self_attn_mask), dim=1
                  )
              if self_attn_padding_mask is not None:
                  if encoder_padding_mask is None:
                      assert encoder_out is not None
                      encoder_padding_mask = self_attn_padding_mask.new_zeros(
                          encoder_out.size(1), encoder_out.size(0)
                      )
                  self_attn_padding_mask = torch.cat(
                      (encoder_padding_mask, self_attn_padding_mask), dim=1
                  )
              assert encoder_out is not None
              y = torch.cat((encoder_out, x), dim=0)
          else:
              y = x
          # print('???x',x)
          # print('???y',y)
          # assert 1==0
          x, attn = self.self_attn(
              query=x,
              key=y,
              value=y,
              key_padding_mask=self_attn_padding_mask,
              incremental_state=incremental_state,
              need_weights=False,
              attn_mask=self_attn_mask,
          )
          x = self.dropout_module(x)
          x = residual + x
          if not self.normalize_before:
              x = self.self_attn_layer_norm(x)
  
          if self.encoder_attn is not None:
              residual = x
              #residual2 = x
              if self.normalize_before:
                  x = self.encoder_attn_layer_norm(x)
              if prev_attn_state is not None:
                  prev_key, prev_value = prev_attn_state[:2]
                  saved_state: Dict[str, Optional[Tensor]] = {
                      "prev_key": prev_key,
                      "prev_value": prev_value,
                  }
                  if len(prev_attn_state) >= 3:
                      saved_state["prev_key_padding_mask"] = prev_attn_state[2]
                  assert incremental_state is not None
                  self.encoder_attn._set_input_buffer(incremental_state, saved_state)
              #print('!!!',x.shape,encoder_out.shape)
              x, attn = self.encoder_attn(
                  query=x,
                  key=encoder_out,
                  value=encoder_out,
                  key_padding_mask=encoder_padding_mask,
                  incremental_state=incremental_state,
                  static_kv=True,
                  need_weights=need_attn or (not self.training and self.need_attn),
                  need_head_weights=need_head_weights,
              )
              x = self.dropout_module(x)
              x = residual + x
              if not self.normalize_before:
                  x = self.encoder_attn_layer_norm(x)
  
          residual = x
          if self.normalize_before:
              x = self.final_layer_norm(x)
  
          x = self.activation_fn(self.fc1(x))
          x = self.activation_dropout_module(x)
          x = self.fc2(x)
          x = self.dropout_module(x)
  
          #if keep_residual==0:
          x = residual + x
          # elif keep_residual==1:
          #     x = residual2 + x
  
          if not self.normalize_before:
              x = self.final_layer_norm(x)
          if self.onnx_trace and incremental_state is not None:
              saved_state = self.self_attn._get_input_buffer(incremental_state)
              assert saved_state is not None
              if self_attn_padding_mask is not None:
                  self_attn_state = [
                      saved_state["prev_key"],
                      saved_state["prev_value"],
                      saved_state["prev_key_padding_mask"],
                  ]
              else:
                  self_attn_state = [saved_state["prev_key"], saved_state["prev_value"]]
              return x, attn, self_attn_state
          return x, attn, None
  
      def make_generation_fast_(self, need_attn: bool = False, **kwargs):
          self.need_attn = need_attn
  
  
  
  class TransformerSentenceEncoderLayer(nn.Module):
      """
      Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
      models.
      """
  
      def __init__(
          self,
          embedding_dim: int = 768,
          ffn_embedding_dim: int = 3072,
          num_attention_heads: int = 8,
          dropout: float = 0.1,
          attention_dropout: float = 0.1,
          activation_dropout: float = 0.1,
          activation_fn: str = 'relu',
          export: bool = False,
          q_noise: float = 0.0,
          qn_block_size: int = 8,
          init_fn: Callable = None,
          scaling_factor: int = 1,
      ) -> None:
          super().__init__()
  
          if init_fn is not None:
              init_fn()
  
          # Initialize parameters
          self.embedding_dim = embedding_dim
          self.dropout_module = FairseqDropout(dropout, module_name=self.__class__.__name__)
          self.activation_dropout_module = FairseqDropout(activation_dropout, module_name=self.__class__.__name__)
  
          # Initialize blocks
          self.activation_fn = get_activation_fn(activation_fn)
          self.self_attn = self.build_self_attention(
              self.embedding_dim,
              num_attention_heads,
              dropout=attention_dropout,
              self_attention=True,
              q_noise=q_noise,
              qn_block_size=qn_block_size,
          )
          # self.self_attn = SelfMultiheadAttention(
          #     self.embedding_dim,
          #     num_attention_heads,
          #     dropout=attention_dropout,
          #     scaling_factor=scaling_factor,
          # )
  
          # layer norm associated with the self attention layer
          self.self_attn_layer_norm = LayerNorm(self.embedding_dim, export=export)
  
          self.fc1 = self.build_fc1(
              self.embedding_dim,
              ffn_embedding_dim,
              q_noise=q_noise,
              qn_block_size=qn_block_size,
          )
          self.fc2 = self.build_fc2(
              ffn_embedding_dim,
              self.embedding_dim,
              q_noise=q_noise,
              qn_block_size=qn_block_size,
          )
  
          # layer norm associated with the position wise feed-forward NN
          self.final_layer_norm = LayerNorm(self.embedding_dim, export=export)
  
      def build_fc1(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(
              nn.Linear(input_dim, output_dim), q_noise, qn_block_size
          )
  
      def build_fc2(self, input_dim, output_dim, q_noise, qn_block_size):
          return quant_noise(
              nn.Linear(input_dim, output_dim), q_noise, qn_block_size
          )
  
      def build_self_attention(
          self,
          embed_dim,
          num_attention_heads,
          dropout,
          self_attention,
          q_noise,
          qn_block_size,
      ):
          return MultiheadAttention(
              embed_dim,
              num_attention_heads,
              dropout=dropout,
              self_attention=True,
              q_noise=q_noise,
              qn_block_size=qn_block_size,
          )
  
      def forward(
          self,
          x: torch.Tensor,
          self_attn_mask: Optional[torch.Tensor] = None,
          self_attn_padding_mask: Optional[torch.Tensor] = None,
      ):
          """
          LayerNorm is applied either before or after the self-attention/ffn
          modules similar to the original Transformer implementation.
          """
          residual = x
          x, attn = self.self_attn(
              query=x,
              key=x,
              value=x,
              key_padding_mask=self_attn_padding_mask,
              need_weights=False,
              attn_mask=self_attn_mask,
          )
          x = self.dropout_module(x)
          x = residual + x
          x = self.self_attn_layer_norm(x)
  
          residual = x
          x = self.activation_fn(self.fc1(x))
          x = self.activation_dropout_module(x)
          x = self.fc2(x)
          x = self.dropout_module(x)
          x = residual + x
          x = self.final_layer_norm(x)
          return x, attn
  
  
  
  def quant_noise(module, p, block_size):
      """
      Wraps modules and applies quantization noise to the weights for
      subsequent quantization with Iterative Product Quantization as
      described in "Training with Quantization Noise for Extreme Model Compression"
  
      Args:
          - module: nn.Module
          - p: amount of Quantization Noise
          - block_size: size of the blocks for subsequent quantization with iPQ
  
      Remarks:
          - Module weights must have the right sizes wrt the block size
          - Only Linear, Embedding and Conv2d modules are supported for the moment
          - For more detail on how to quantize by blocks with convolutional weights,
            see "And the Bit Goes Down: Revisiting the Quantization of Neural Networks"
          - We implement the simplest form of noise here as stated in the paper
            which consists in randomly dropping blocks
      """
  
      # if no quantization noise, don't register hook
      if p <= 0:
          return module
  
      # supported modules
      assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d))
  
      # test whether module.weight has the right sizes wrt block_size
      is_conv = module.weight.ndim == 4
  
      # 2D matrix
      if not is_conv:
          assert module.weight.size(1) % block_size == 0, "Input features must be a multiple of block sizes"
  
      # 4D matrix
      else:
          # 1x1 convolutions
          if module.kernel_size == (1, 1):
              assert module.in_channels % block_size == 0, "Input channels must be a multiple of block sizes"
          # regular convolutions
          else:
              k = module.kernel_size[0] * module.kernel_size[1]
              assert k % block_size == 0, "Kernel size must be a multiple of block size"
  
      def _forward_pre_hook(mod, input):
          # no noise for evaluation
          if mod.training:
              if not is_conv:
                  # gather weight and sizes
                  weight = mod.weight
                  in_features = weight.size(1)
                  out_features = weight.size(0)
  
                  # split weight matrix into blocks and randomly drop selected blocks
                  mask = torch.zeros(in_features // block_size * out_features, device=weight.device)
                  mask.bernoulli_(p)
                  mask = mask.repeat_interleave(block_size, -1).view(-1, in_features)
  
              else:
                  # gather weight and sizes
                  weight = mod.weight
                  in_channels = mod.in_channels
                  out_channels = mod.out_channels
  
                  # split weight matrix into blocks and randomly drop selected blocks
                  if mod.kernel_size == (1, 1):
                      mask = torch.zeros(int(in_channels // block_size * out_channels), device=weight.device)
                      mask.bernoulli_(p)
                      mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels)
                  else:
                      mask = torch.zeros(weight.size(0), weight.size(1), device=weight.device)
                      mask.bernoulli_(p)
                      mask = mask.unsqueeze(2).unsqueeze(3).repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])
  
              # scale weights and apply mask
              mask = mask.to(torch.bool)  # x.bool() is not currently supported in TorchScript
              s = 1 / (1 - p)
              mod.weight.data = s * weight.masked_fill(mask, 0)
  
      module.register_forward_pre_hook(_forward_pre_hook)
      return module
  
  
  
  
  
  