text: import math
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  class PositionalEncoding(nn.Module):
  
      def __init__(self, dropout, dim, max_len=5000):
  
          pe = torch.zeros(max_len, dim)
          div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *
                               -(math.log(10000.0) / dim)))
          position = torch.arange(0, max_len).unsqueeze(1)
          pe[:, 0::2] = torch.sin(position.float() * div_term)
          pe[:, 1::2] = torch.cos(position.float() * div_term)
          pe = pe.unsqueeze(0)
  
          super(PositionalEncoding, self).__init__()
          self.register_buffer('pe', pe)
          self.dropout = nn.Dropout(p=dropout)
          self.dim = dim
  
      def forward(self, emb, step=None, add_emb=None):
          emb = emb * math.sqrt(self.dim)
          if add_emb is not None:
              emb = emb + add_emb
          if (step):
              pos = self.pe[:, step][:, None, :]
              emb = emb + pos
          else:
              pos = self.pe[:, :emb.size(1)]
              emb = emb + pos
          emb = self.dropout(emb)
          return emb
  
  
  def gumbel_softmax(logits, tau=1.0, hard=False, log_mode=True, dim=-1):
  
      while(True):
          gumbels = -torch.empty_like(logits).exponential_().log()  # ~Gumbel(0,1)
          gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
          if log_mode:
              y_soft = gumbels.log_softmax(dim)
          else:
              y_soft = gumbels.softmax(dim)
          if torch.sum(torch.isnan(y_soft)).item() < 0.01:
              break
  
      if hard:
          # Straight through.
          index = y_soft.max(dim, keepdim=True)[1]
          y_hard = torch.zeros_like(logits).scatter_(dim, index, 1.0)
          ret = y_hard - y_soft.detach() + y_soft
      else:
          # Reparametrization trick.
          ret = y_soft
      return ret
  
  
  class MultiHeadedAttention(nn.Module):
      """
      Multi-Head Attention module from
      "Attention is All You Need"
      :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.
      Similar to standard `dot` attention but uses
      multiple attention distributions simulataneously
      to select relevant items.
      .. mermaid::
         graph BT
            A[key]
            B[value]
            C[query]
            O[output]
            subgraph Attn
              D[Attn 1]
              E[Attn 2]
              F[Attn N]
            end
            A --> D
            C --> D
            A --> E
            C --> E
            A --> F
            C --> F
            D --> O
            E --> O
            F --> O
            B --> O
      Also includes several additional tricks.
      Args:
         head_count (int): number of parallel heads
         model_dim (int): the dimension of keys/values/queries,
             must be divisible by head_count
         dropout (float): dropout parameter
      """
  
      def __init__(self, head_count, model_dim, dropout=0.1, use_final_linear=True):
          assert model_dim % head_count == 0
          self.dim_per_head = model_dim // head_count
          self.model_dim = model_dim
  
          super(MultiHeadedAttention, self).__init__()
          self.head_count = head_count
  
          self.linear_keys = nn.Linear(model_dim,
                                       head_count * self.dim_per_head)
          self.linear_values = nn.Linear(model_dim,
                                         head_count * self.dim_per_head)
          self.linear_query = nn.Linear(model_dim,
                                        head_count * self.dim_per_head)
          self.softmax = nn.Softmax(dim=-1)
          self.dropout = nn.Dropout(dropout)
          self.use_final_linear = use_final_linear
          if (self.use_final_linear):
              self.final_linear = nn.Linear(model_dim, model_dim)
  
      def forward(self, key, value, query, mask=None, layer_cache=None,
                  type=None, predefined_graph_1=None, requires_att=False):
          """
          Compute the context vector and the attention vectors.
          Args:
             key (`FloatTensor`): set of `key_len`
                  key vectors `[batch, key_len, dim]`
             value (`FloatTensor`): set of `key_len`
                  value vectors `[batch, key_len, dim]`
             query (`FloatTensor`): set of `query_len`
                   query vectors  `[batch, query_len, dim]`
             mask: binary mask indicating which keys have
                   non-zero attention `[batch, query_len, key_len]`
          Returns:
             (`FloatTensor`, `FloatTensor`) :
             * output context vectors `[batch, query_len, dim]`
             * one of the attention vectors `[batch, query_len, key_len]`
          """
  
          batch_size = key.size(0)
          dim_per_head = self.dim_per_head
          head_count = self.head_count
  
          def shape(x):
              """  projection """
              return x.view(batch_size, -1, head_count, dim_per_head) \
                  .transpose(1, 2)
  
          def unshape(x):
              """  compute context """
              return x.transpose(1, 2).contiguous() \
                  .view(batch_size, -1, head_count * dim_per_head)
  
          # 1) Project key, value, and query.
          if layer_cache is not None:
              if type == "self":
                  query, key, value = self.linear_query(query), \
                                      self.linear_keys(query), \
                                      self.linear_values(query)
  
                  key = shape(key)
                  value = shape(value)
  
                  if layer_cache is not None:
                      device = key.device
                      if layer_cache["self_keys"] is not None:
                          key = torch.cat(
                              (layer_cache["self_keys"].to(device), key),
                              dim=2)
                      if layer_cache["self_values"] is not None:
                          value = torch.cat(
                              (layer_cache["self_values"].to(device), value),
                              dim=2)
                      layer_cache["self_keys"] = key
                      layer_cache["self_values"] = value
              elif type == "context":
                  query = self.linear_query(query)
                  if layer_cache is not None:
                      if layer_cache["memory_keys"] is None:
                          key, value = self.linear_keys(key), \
                                       self.linear_values(value)
                          key = shape(key)
                          value = shape(value)
                      else:
                          key, value = layer_cache["memory_keys"], \
                                       layer_cache["memory_values"]
                      layer_cache["memory_keys"] = key
                      layer_cache["memory_values"] = value
                  else:
                      key, value = self.linear_keys(key), \
                                   self.linear_values(value)
                      key = shape(key)
                      value = shape(value)
          else:
              key = self.linear_keys(key)
              value = self.linear_values(value)
              query = self.linear_query(query)
              key = shape(key)
              value = shape(value)
  
          query = shape(query)
  
          # 2) Calculate and scale scores.
          query = query / math.sqrt(dim_per_head)
          scores = torch.matmul(query, key.transpose(2, 3))
  
          if mask is not None:
              mask = mask.unsqueeze(1).expand_as(scores)
              scores = scores.masked_fill(mask, -1e4)
  
          # 3) Apply attention dropout and compute context vectors.
  
          attn = self.softmax(scores)
  
          if requires_att:
              required_att = attn.mean(1)
          else:
              required_att = None
  
          if (predefined_graph_1 is not None):
              attn_masked = attn[:, -1] * predefined_graph_1
              attn_masked = attn_masked / (torch.sum(attn_masked, 2).unsqueeze(2) + 1e-9)
  
              attn = torch.cat([attn[:, :-1], attn_masked.unsqueeze(1)], 1)
  
          drop_attn = self.dropout(attn)
          if (self.use_final_linear):
              context = unshape(torch.matmul(drop_attn, value))
              output = self.final_linear(context)
              return output, required_att
          else:
              context = torch.matmul(drop_attn, value)
              return context, required_att
  
  
  class DecoderState(object):
      """Interface for grouping together the current state of a recurrent
      decoder. In the simplest case just represents the hidden state of
      the model.  But can also be used for implementing various forms of
      input_feeding and non-recurrent models.
      Modules need to implement this to utilize beam search decoding.
      """
      def detach(self):
          """ Need to document this """
          self.hidden = tuple([_.detach() for _ in self.hidden])
          self.input_feed = self.input_feed.detach()
  
      def beam_update(self, idx, positions, beam_size):
          """ Need to document this """
          for e in self._all:
              sizes = e.size()
              br = sizes[1]
              if len(sizes) == 3:
                  sent_states = e.view(sizes[0], beam_size, br // beam_size,
                                       sizes[2])[:, :, idx]
              else:
                  sent_states = e.view(sizes[0], beam_size,
                                       br // beam_size,
                                       sizes[2],
                                       sizes[3])[:, :, idx]
  
              sent_states.data.copy_(
                  sent_states.data.index_select(1, positions))
  
      def map_batch_fn(self, fn):
          raise NotImplementedError()
  
  
  def gelu(x):
      return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
  
  
  class PositionwiseFeedForward(nn.Module):
      """ A two-layer Feed-Forward-Network with residual layer norm.
      Args:
          d_model (int): the size of input for the first-layer of the FFN.
          d_ff (int): the hidden layer size of the second-layer
              of the FNN.
          dropout (float): dropout probability in :math:`[0, 1)`.
      """
  
      def __init__(self, d_model, d_ff, dropout=0.1):
          super(PositionwiseFeedForward, self).__init__()
          self.w_1 = nn.Linear(d_model, d_ff)
          self.w_2 = nn.Linear(d_ff, d_model)
          self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
          self.actv = gelu
          self.dropout_1 = nn.Dropout(dropout)
          self.dropout_2 = nn.Dropout(dropout)
  
      def forward(self, x):
          inter = self.dropout_1(self.actv(self.w_1(self.layer_norm(x))))
          output = self.dropout_2(self.w_2(inter))
          return output + x
  