text: #!/usr/bin/env python3
  # Copyright 2018-present, HKUST-KnowComp.
  # All rights reserved.
  #
  # This source code is licensed under the license found in the
  # LICENSE file in the root directory of this source tree.
  """Definitions of model layers/NN modules"""
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.autograd import Variable
  import math
  import random
  
  
  # ------------------------------------------------------------------------------
  # Modules
  # ------------------------------------------------------------------------------
  
  
  class StackedBRNN(nn.Module):
      """Stacked Bi-directional RNNs.
  
      Differs from standard PyTorch library in that it has the option to save
      and concat the hidden states between layers. (i.e. the output hidden size
      for each sequence input is num_layers * hidden_size).
      """
  
      def __init__(self, input_size, hidden_size, num_layers,
                   dropout_rate=0, dropout_output=False, rnn_type=nn.LSTM,
                   concat_layers=False, padding=False):
          super(StackedBRNN, self).__init__()
          self.padding = padding
          self.dropout_output = dropout_output
          self.dropout_rate = dropout_rate
          self.num_layers = num_layers
          self.concat_layers = concat_layers
          self.rnns = nn.ModuleList()
          for i in range(num_layers):
              input_size = input_size if i == 0 else 2 * hidden_size
              self.rnns.append(rnn_type(input_size, hidden_size,
                                        num_layers=1,
                                        bidirectional=True))
  
      def forward(self, x, x_mask):
          """Encode either padded or non-padded sequences.
  
          Can choose to either handle or ignore variable length sequences.
          Always handle padding in eval.
  
          Args:
              x: batch * len * hdim
              x_mask: batch * len (1 for padding, 0 for true)
          Output:
              x_encoded: batch * len * hdim_encoded
          """
          if x_mask.data.sum() == 0:
              # No padding necessary.
              output = self._forward_unpadded(x, x_mask)
          elif self.padding or not self.training:
              # Pad if we care or if its during eval.
              output = self._forward_padded(x, x_mask)
          else:
              # We don't care.
              output = self._forward_unpadded(x, x_mask)
  
          return output.contiguous()
  
      def _forward_unpadded(self, x, x_mask):
          """Faster encoding that ignores any padding."""
          # Transpose batch and sequence dims
          x = x.transpose(0, 1)
  
          # Encode all layers
          outputs = [x]
          for i in range(self.num_layers):
              rnn_input = outputs[-1]
  
              # Apply dropout to hidden input
              if self.dropout_rate > 0:
                  rnn_input = F.dropout(rnn_input,
                                        p=self.dropout_rate,
                                        training=self.training)
              # Forward
              rnn_output = self.rnns[i](rnn_input)[0]
              outputs.append(rnn_output)
  
          # Concat hidden layers
          if self.concat_layers:
              output = torch.cat(outputs[1:], 2)
          else:
              output = outputs[-1]
  
          # Transpose back
          output = output.transpose(0, 1)
  
          # Dropout on output layer
          if self.dropout_output and self.dropout_rate > 0:
              output = F.dropout(output,
                                 p=self.dropout_rate,
                                 training=self.training)
          return output
  
      def _forward_padded(self, x, x_mask):
          """Slower (significantly), but more precise, encoding that handles
          padding.
          """
          # Compute sorted sequence lengths
          lengths = x_mask.data.eq(0).long().sum(1).squeeze()
          _, idx_sort = torch.sort(lengths, dim=0, descending=True)
          _, idx_unsort = torch.sort(idx_sort, dim=0)
  
          lengths = list(lengths[idx_sort])
          idx_sort = Variable(idx_sort)
          idx_unsort = Variable(idx_unsort)
  
          # Sort x
          x = x.index_select(0, idx_sort)
  
          # Transpose batch and sequence dims
          x = x.transpose(0, 1)
  
          # Pack it up
          rnn_input = nn.utils.rnn.pack_padded_sequence(x, lengths)
  
          # Encode all layers
          outputs = [rnn_input]
          for i in range(self.num_layers):
              rnn_input = outputs[-1]
  
              # Apply dropout to input
              if self.dropout_rate > 0:
                  dropout_input = F.dropout(rnn_input.data,
                                            p=self.dropout_rate,
                                            training=self.training)
                  rnn_input = nn.utils.rnn.PackedSequence(dropout_input,
                                                          rnn_input.batch_sizes)
              outputs.append(self.rnns[i](rnn_input)[0])
  
          # Unpack everything
          for i, o in enumerate(outputs[1:], 1):
              outputs[i] = nn.utils.rnn.pad_packed_sequence(o)[0]
  
          # Concat hidden layers or take final
          if self.concat_layers:
              output = torch.cat(outputs[1:], 2)
          else:
              output = outputs[-1]
  
          # Transpose and unsort
          output = output.transpose(0, 1)
          output = output.index_select(0, idx_unsort)
  
          # Pad up to original batch sequence length
          if output.size(1) != x_mask.size(1):
              padding = torch.zeros(output.size(0),
                                    x_mask.size(1) - output.size(1),
                                    output.size(2)).type(output.data.type())
              output = torch.cat([output, Variable(padding)], 1)
  
          # Dropout on output layer
          if self.dropout_output and self.dropout_rate > 0:
              output = F.dropout(output,
                                 p=self.dropout_rate,
                                 training=self.training)
          return output
  
  class MaxoutRNN(nn.Module):
      """Maxout Bi-directional RNNs.
  
      Differs from standard PyTorch library in that it has the option to save
      and concat the hidden states between layers. (i.e. the output hidden size
      for each sequence input is num_layers * hidden_size).
      """
  
      def __init__(self, input_size, hidden_size, num_layers,
                   dropout_rate=0, dropout_output=False, rnn_type=nn.LSTM,
                   padding=False):
          super(MaxoutRNN, self).__init__()
          self.padding = padding
          self.dropout_output = dropout_output
          self.dropout_rate = dropout_rate
          self.num_layers = num_layers
          self.rnns = nn.ModuleList()
          for i in range(num_layers):
              input_size = input_size if i == 0 else 2 * hidden_size
              self.rnns.append(rnn_type(input_size, hidden_size,
                                        num_layers=1,
                                        bidirectional=True))
  
      def forward(self, x, x_mask):
          """Encode either padded or non-padded sequences.
  
          Can choose to either handle or ignore variable length sequences.
          Always handle padding in eval.
  
          Args:
              x: batch * len * hdim
              x_mask: batch * len (1 for padding, 0 for true)
          Output:
              x_encoded: batch * len * hdim_encoded
          """
          if x_mask.data.sum() == 0:
              # No padding necessary.
              output = self._forward_unpadded(x, x_mask)
          elif self.padding or not self.training:
              # Pad if we care or if its during eval.
              output = self._forward_padded(x, x_mask)
          else:
              # We don't care.
              output = self._forward_unpadded(x, x_mask)
  
          return output.contiguous()
  
      def _forward_unpadded(self, x, x_mask):
          """Faster encoding that ignores any padding."""
          # Transpose batch and sequence dims
          x = x.transpose(0, 1)
  
          # Encode all layers
          outputs = [x]
          for i in range(self.num_layers):
              rnn_input = outputs[-1]
  
              # Apply dropout to hidden input
              if self.dropout_rate > 0:
                  rnn_input = F.dropout(rnn_input,
                                        p=self.dropout_rate,
                                        training=self.training)
              # Forward
              rnn_output = self.rnns[i](rnn_input)[0]
              outputs.append(rnn_output)
  
          # Maxout hidden layers
          max_output = outputs[1]
          for i in range(2, self.num_layers):
              max_output = torch.max(max_output, outputs[i])
  
          # Transpose back
          output = max_output.transpose(0, 1)
  
          # Dropout on output layer
          if self.dropout_output and self.dropout_rate > 0:
              output = F.dropout(output,
                                 p=self.dropout_rate,
                                 training=self.training)
          return output
  
      def _forward_padded(self, x, x_mask):
          """Slower (significantly), but more precise, encoding that handles
          padding.
          """
          # Compute sorted sequence lengths
          lengths = x_mask.data.eq(0).long().sum(1).squeeze()
          _, idx_sort = torch.sort(lengths, dim=0, descending=True)
          _, idx_unsort = torch.sort(idx_sort, dim=0)
  
          lengths = list(lengths[idx_sort])
          idx_sort = Variable(idx_sort)
          idx_unsort = Variable(idx_unsort)
  
          # Sort x
          x = x.index_select(0, idx_sort)
  
          # Transpose batch and sequence dims
          x = x.transpose(0, 1)
  
          # Pack it up
          rnn_input = nn.utils.rnn.pack_padded_sequence(x, lengths)
  
          # Encode all layers
          outputs = [rnn_input]
          for i in range(self.num_layers):
              rnn_input = outputs[-1]
  
              # Apply dropout to input
              if self.dropout_rate > 0:
                  dropout_input = F.dropout(rnn_input.data,
                                            p=self.dropout_rate,
                                            training=self.training)
                  rnn_input = nn.utils.rnn.PackedSequence(dropout_input,
                                                          rnn_input.batch_sizes)
              outputs.append(self.rnns[i](rnn_input)[0])
  
          # Unpack everything
          for i, o in enumerate(outputs[1:], 1):
              outputs[i] = nn.utils.rnn.pad_packed_sequence(o)[0]
  
          # Maxout hidden layers
          max_output = outputs[1]
          for i in range(2, self.num_layers):
              max_output = torch.max(max_output, outputs[i])
  
          # Transpose and unsort
          output = max_output.transpose(0, 1)
          output = output.index_select(0, idx_unsort)
  
          # Pad up to original batch sequence length
          if output.size(1) != x_mask.size(1):
              padding = torch.zeros(output.size(0),
                                    x_mask.size(1) - output.size(1),
                                    output.size(2)).type(output.data.type())
              output = torch.cat([output, Variable(padding)], 1)
  
          # Dropout on output layer
          if self.dropout_output and self.dropout_rate > 0:
              output = F.dropout(output,
                                 p=self.dropout_rate,
                                 training=self.training)
          return output
  
  
  class GatedAttnRNN(nn.Module):
      """Gated Attention-baed RNN.
          Given sequences X and Y, match sequence Y to each element in X.
      """
      def __init__(self, x_size, y_size, hidden_size, dropout_rate=0, cell_type=nn.GRUCell):
          super(GatedAttnRNN, self).__init__()
          self.hidden_size = hidden_size
          self.dropout_rate = dropout_rate
          self.cell = cell_type(x_size + y_size, hidden_size)
          self.linear = nn.Linear(x_size + y_size + hidden_size, hidden_size, bias=False)
          self.weight = nn.Linear(hidden_size, 1, bias=False)
          self.gate = Gate(x_size + y_size)
      
      def forward(self, x, y, y_mask):
          """
          Args:
              x: batch * len1 * dim1
              y: batch * len2 * dim2
              y_mask: batch * len2 (1 for padding, 0 for true)
          Output:
              matched_seq: batch * len1 * hdim
          """
          v0 = Variable(torch.zeros([x.size(0), self.hidden_size]))
          if y.data.is_cuda:
              v0 = v0.cuda()
          output = [v0.unsqueeze(1)]
          x_len = x.size(1)
          for i in range(x_len):
              u = x[:, i, :]
              v = output[-1]
              s = F.tanh(self.linear(torch.cat([y, u.unsqueeze(1).repeat(1, y.size(1), 1), v.repeat(1, y.size(1), 1)], 2)))
              w_s = self.weight(s).squeeze(2)
              w_s.data.masked_fill_(y_mask.data, -float('inf'))
              a = F.softmax(w_s)
              c = a.unsqueeze(1).bmm(y).squeeze(1)
              u_c = self.gate(torch.cat([u, c], 1))
              v_t = self.cell(u_c, v.squeeze(1))
              output.append(v_t.unsqueeze(1))
          res = torch.cat(output[1:], 1)
          return F.dropout(res, p=self.dropout_rate, training=self.training)
  
  class FeedForwardNetwork(nn.Module):
      def __init__(self, input_size, hidden_size, output_size, dropout_rate=0):
          super(FeedForwardNetwork, self).__init__()
          self.dropout_rate = dropout_rate
          self.linear1 = nn.Linear(input_size, hidden_size)
          self.linear2 = nn.Linear(hidden_size, output_size)
  
      def forward(self, x):
          x_proj = F.dropout(F.relu(self.linear1(x)), p=self.dropout_rate, training=self.training)
          x_proj = self.linear2(x_proj)
          return x_proj
              
  
  class PointerNetwork(nn.Module):
      def __init__(self, x_size, y_size, hidden_size, dropout_rate=0, cell_type=nn.GRUCell, normalize=True):
          super(PointerNetwork, self).__init__()
          self.normalize = normalize
          self.hidden_size = hidden_size
          self.dropout_rate = dropout_rate
          self.linear = nn.Linear(x_size+y_size, hidden_size, bias=False)
          self.weights = nn.Linear(hidden_size, 1, bias=False)
          self.self_attn = NonLinearSeqAttn(y_size, hidden_size)
          self.cell = cell_type(x_size, y_size)
  
      def init_hiddens(self, y, y_mask):
          attn = self.self_attn(y, y_mask)
          res = attn.unsqueeze(1).bmm(y).squeeze(1) # [B, I]
          return res
      
      def pointer(self, x, state, x_mask):
          x_ = torch.cat([x, state.unsqueeze(1).repeat(1,x.size(1),1)], 2)
          s0 = F.tanh(self.linear(x_))
          s = self.weights(s0).view(x.size(0), x.size(1))
          s.data.masked_fill_(x_mask.data, -float('inf'))
          a = F.softmax(s)
          res = a.unsqueeze(1).bmm(x).squeeze(1)
          if self.normalize:
              if self.training:
                  # In training we output log-softmax for NLL
                  scores = F.log_softmax(s)
              else:
                  # ...Otherwise 0-1 probabilities
                  scores = F.softmax(s)
          else:
              scores = a.exp()
          return res, scores
  
  
      def forward(self, x, y, x_mask, y_mask):
          hiddens = self.init_hiddens(y, y_mask)
          c, start_scores = self.pointer(x, hiddens, x_mask)
          c_ = F.dropout(c, p=self.dropout_rate, training=self.training)
          hiddens = self.cell(c_, hiddens)
          c, end_scores = self.pointer(x, hiddens, x_mask)
          return start_scores, end_scores
  
  class MemoryAnsPointer(nn.Module):
      def __init__(self, x_size, y_size, hidden_size, hop=1, dropout_rate=0, normalize=True):
          super(MemoryAnsPointer, self).__init__()
          self.normalize = normalize
          self.hidden_size = hidden_size
          self.hop = hop
          self.dropout_rate = dropout_rate
          self.FFNs_start = nn.ModuleList()
          self.SFUs_start = nn.ModuleList()
          self.FFNs_end = nn.ModuleList()
          self.SFUs_end = nn.ModuleList()
          for i in range(self.hop):
              self.FFNs_start.append(FeedForwardNetwork(x_size+y_size+2*hidden_size, hidden_size, 1, dropout_rate))
              self.SFUs_start.append(SFU(y_size, 2*hidden_size))
              self.FFNs_end.append(FeedForwardNetwork(x_size+y_size+2*hidden_size, hidden_size, 1, dropout_rate))
              self.SFUs_end.append(SFU(y_size, 2*hidden_size))
      
      def forward(self, x, y, x_mask, y_mask):
          z_s = y[:,-1,:].unsqueeze(1) # [B, 1, I]
          z_e = None
          s = None
          e = None
          p_s = None
          p_e = None
          
          for i in range(self.hop):
              z_s_ = z_s.repeat(1,x.size(1),1) # [B, S, I]
              s = self.FFNs_start[i](torch.cat([x, z_s_, x*z_s_], 2)).squeeze(2)
              s.data.masked_fill_(x_mask.data, -float('inf'))
              p_s = F.softmax(s, dim=1) # [B, S]
              u_s = p_s.unsqueeze(1).bmm(x) # [B, 1, I]
              z_e = self.SFUs_start[i](z_s, u_s) # [B, 1, I]
              z_e_ = z_e.repeat(1,x.size(1),1) # [B, S, I]
              e = self.FFNs_end[i](torch.cat([x, z_e_, x*z_e_], 2)).squeeze(2)
              e.data.masked_fill_(x_mask.data, -float('inf'))
              p_e = F.softmax(e, dim=1) # [B, S]
              u_e = p_e.unsqueeze(1).bmm(x) # [B, 1, I]
              z_s = self.SFUs_end[i](z_e, u_e)
          if self.normalize:
              if self.training:
                  # In training we output log-softmax for NLL
                  p_s = F.log_softmax(s, dim=1) # [B, S]
                  p_e = F.log_softmax(e, dim=1) # [B, S]
              else:
                  # ...Otherwise 0-1 probabilities
                  p_s = F.softmax(s, dim=1) # [B, S]
                  p_e = F.softmax(e, dim=1) # [B, S]
          else:
              p_s = s.exp()
              p_e = e.exp()
          return p_s, p_e
  
  class StochasticAnsPointer(nn.Module):
      def __init__(self, x_size, y_size, hidden_size, hop=1, dropout_rate=0, normalize=True):
          super(StochasticAnsPointer, self).__init__()
          """"StochasticAnsPointer
          """
          self.normalize = normalize
          self.hidden_size = hidden_size
          self.hop = hop
          self.dropout_rate = dropout_rate
          # self.self_attn = LinearSeqAttn(y_size)
          self.self_attn = NonLinearSeqAttn(y_size, hidden_size)
          self.FFNs_start = nn.ModuleList()
          self.SFUs_start = nn.ModuleList()
          self.FFNs_end = nn.ModuleList()
          self.SFUs_end = nn.ModuleList()
          for i in range(self.hop):
              self.FFNs_start.append(FeedForwardNetwork(x_size+y_size+2*hidden_size, hidden_size, 1, dropout_rate))
              self.SFUs_start.append(SFU(y_size, 2*hidden_size))
              self.FFNs_end.append(FeedForwardNetwork(x_size+y_size+2*hidden_size, hidden_size, 1, dropout_rate))
              self.SFUs_end.append(SFU(y_size, 2*hidden_size))
          """Bilinear Term
          """
          self.self_attn_linear = LinearSeqAttn(y_size)
          self.start_attn_liner = nn.Linear(y_size, x_size)
          self.end_attn_liner = nn.Linear(y_size, x_size)
  
      def forward(self, x, y, x_mask, y_mask):
          """"MemoryAnsPointer
          """
          attn = self.self_attn(y, y_mask)
          z_s = attn.unsqueeze(1).bmm(y) # [B, 1, I]
          z_e = None
          s = None
          e = None
          p_s = None
          p_e = None
          s_list = []
          e_list = []
          
          for i in range(self.hop):
              z_s_ = z_s.repeat(1,x.size(1),1) # [B, S, I]
              s = self.FFNs_start[i](torch.cat([x, z_s_, x*z_s_], 2)).squeeze(2)
              s.data.masked_fill_(x_mask.data, -float('inf'))
              s_list.append(s)
              p_s = F.softmax(s, dim=1) # [B, S]
              u_s = p_s.unsqueeze(1).bmm(x) # [B, 1, I]
              z_e = self.SFUs_start[i](z_s, u_s) # [B, 1, I]
              z_e_ = z_e.repeat(1,x.size(1),1) # [B, S, I]
              e = self.FFNs_end[i](torch.cat([x, z_e_, x*z_e_], 2)).squeeze(2)
              e.data.masked_fill_(x_mask.data, -float('inf'))
              e_list.append(e)
              p_e = F.softmax(e, dim=1) # [B, S]
              u_e = p_e.unsqueeze(1).bmm(x) # [B, 1, I]
              z_s = self.SFUs_end[i](z_e, u_e)
          """Bilinear Term
          """
          weights = self.self_attn_linear(y, y_mask)
          attn = weights.unsqueeze(1).bmm(y).squeeze(1)
          start_Wy = self.start_attn_liner(attn)
          start_xWy = x.bmm(start_Wy.unsqueeze(2)).squeeze(2)
          start_xWy.data.masked_fill_(x_mask.data, -float('inf'))
          end_Wy = self.end_attn_liner(attn)
          end_xWy = x.bmm(end_Wy.unsqueeze(2)).squeeze(2)
          end_xWy.data.masked_fill_(x_mask.data, -float('inf'))
  
          # Merge two results
          s_list_ = [start_xWy.unsqueeze(2)]
          e_list_ = [end_xWy.unsqueeze(2)]
          if self.training:
              # In training we output log-softmax for NLL
              for s in s_list:
                  if random.random() < self.dropout_rate:
                      continue
                  s_list_.append(s.unsqueeze(2))
              for e in e_list:
                  if random.random() < self.dropout_rate:
                      continue
                  e_list_.append(e.unsqueeze(2))
              s = torch.cat(s_list_, 2).mean(2)
              e = torch.cat(e_list_, 2).mean(2)
              
          else:
              for s in s_list:
                  s_list_.append(s.unsqueeze(2))
              for e in e_list:
                  e_list_.append(e.unsqueeze(2))
              s = torch.cat(s_list_, 2).mean(2)
              e = torch.cat(e_list_, 2).mean(2)
          if self.normalize:
              if self.training:
                  p_s = F.log_softmax(s, dim=1) # [B, S]
                  p_e = F.log_softmax(e, dim=1) # [B, S]
              else:
                  p_s = F.softmax(s, dim=1) # [B, S]
                  p_e = F.softmax(e, dim=1) # [B, S]
          else:
              p_s = s.exp() # [B, S]
              p_e = e.exp() # [B, S]
          return p_s, p_e
  
  
  
  # ------------------------------------------------------------------------------
  # Attentions
  # ------------------------------------------------------------------------------
  
  class SeqAttnMatch(nn.Module):
      """Given sequences X and Y, match sequence Y to each element in X.
  
      * o_i = sum(alpha_j * y_j) for i in X
      * alpha_j = softmax(y_j * x_i)
      """
  
      def __init__(self, input_size, identity=False):
          super(SeqAttnMatch, self).__init__()
          if not identity:
              self.linear = nn.Linear(input_size, input_size)
          else:
              self.linear = None
  
      def forward(self, x, y, y_mask):
          """
          Args:
              x: batch * len1 * hdim
              y: batch * len2 * hdim
              y_mask: batch * len2 (1 for padding, 0 for true)
          Output:
              matched_seq: batch * len1 * hdim
          """
          # Project vectors
          if self.linear:
              x_proj = self.linear(x.view(-1, x.size(2))).view(x.size())
              x_proj = F.relu(x_proj)
              y_proj = self.linear(y.view(-1, y.size(2))).view(y.size())
              y_proj = F.relu(y_proj)
          else:
              x_proj = x
              y_proj = y
  
          # Compute scores
          scores = x_proj.bmm(y_proj.transpose(2, 1))
  
          # Mask padding
          y_mask = y_mask.unsqueeze(1).expand(scores.size())
          scores.data.masked_fill_(y_mask.data, -float('inf'))
  
          # Normalize with softmax
          alpha = F.softmax(scores, dim=2)
  
          # Take weighted average
          matched_seq = alpha.bmm(y)
          return matched_seq
  
  class SelfAttnMatch(nn.Module):
      """Given sequences X and Y, match sequence Y to each element in X.
  
      * o_i = sum(alpha_j * x_j) for i in X
      * alpha_j = softmax(x_j * x_i)
      """
  
      def __init__(self, input_size, identity=False, diag=True):
          super(SelfAttnMatch, self).__init__()
          if not identity:
              self.linear = nn.Linear(input_size, input_size)
          else:
              self.linear = None
          self.diag = diag
  
      def forward(self, x, x_mask):
          """
          Args:
              x: batch * len1 * dim1
              x_mask: batch * len1 (1 for padding, 0 for true)
          Output:
              matched_seq: batch * len1 * dim1
          """
          # Project vectors
          if self.linear:
              x_proj = self.linear(x.view(-1, x.size(2))).view(x.size())
              x_proj = F.relu(x_proj)
          else:
              x_proj = x
  
          # Compute scores
          scores = x_proj.bmm(x_proj.transpose(2, 1))
          if not self.diag:
              x_len = x.size(1)
              for i in range(x_len):
                  scores[:, i, i] = 0
  
          # Mask padding
          x_mask = x_mask.unsqueeze(1).expand(scores.size())
          scores.data.masked_fill_(x_mask.data, -float('inf'))
  
          # Normalize with softmax
          alpha = F.softmax(scores, dim=2)
  
          # Take weighted average
          matched_seq = alpha.bmm(x)
          return matched_seq
  
  
  class BilinearSeqAttn(nn.Module):
      """A bilinear attention layer over a sequence X w.r.t y:
  
      * o_i = softmax(x_i'Wy) for x_i in X.
  
      Optionally don't normalize output weights.
      """
  
      def __init__(self, x_size, y_size, identity=False, normalize=True):
          super(BilinearSeqAttn, self).__init__()
          self.normalize = normalize
  
          # If identity is true, we just use a dot product without transformation.
          if not identity:
              self.linear = nn.Linear(y_size, x_size)
          else:
              self.linear = None
  
      def forward(self, x, y, x_mask):
          """
          Args:
              x: batch * len * hdim1
              y: batch * hdim2
              x_mask: batch * len (1 for padding, 0 for true)
          Output:
              alpha = batch * len
          """
          Wy = self.linear(y) if self.linear is not None else y
          xWy = x.bmm(Wy.unsqueeze(2)).squeeze(2)
          xWy.data.masked_fill_(x_mask.data, -float('inf'))
          if self.normalize:
              if self.training:
                  # In training we output log-softmax for NLL
                  alpha = F.log_softmax(xWy)
              else:
                  # ...Otherwise 0-1 probabilities
                  alpha = F.softmax(xWy)
          else:
              alpha = xWy.exp()
          return alpha
  
  
  class LinearSeqAttn(nn.Module):
      """Self attention over a sequence:
  
      * o_i = softmax(Wx_i) for x_i in X.
      """
  
      def __init__(self, input_size):
          super(LinearSeqAttn, self).__init__()
          self.linear = nn.Linear(input_size, 1)
  
      def forward(self, x, x_mask):
          """
          Args:
              x: batch * len * hdim
              x_mask: batch * len (1 for padding, 0 for true)
          Output:
              alpha: batch * len
          """
          x_flat = x.view(-1, x.size(-1))
          scores = self.linear(x_flat).view(x.size(0), x.size(1))
          scores.data.masked_fill_(x_mask.data, -float('inf'))
          alpha = F.softmax(scores)
          return alpha
  
  class NonLinearSeqAttn(nn.Module):
      """Self attention over a sequence:
  
      * o_i = softmax(function(Wx_i)) for x_i in X.
      """
  
      def __init__(self, input_size, hidden_size):
          super(NonLinearSeqAttn, self).__init__()
          self.FFN = FeedForwardNetwork(input_size, hidden_size, 1)
  
      def forward(self, x, x_mask):
          """
          Args:
              x: batch * len * dim
              x_mask: batch * len (1 for padding, 0 for true)
          Output:
              alpha: batch * len
          """
          scores = self.FFN(x).squeeze(2)
          scores.data.masked_fill_(x_mask.data, -float('inf'))
          alpha = F.softmax(scores)
          return alpha
  
  
  # ------------------------------------------------------------------------------
  # Functional Units
  # ------------------------------------------------------------------------------
  
  class Gate(nn.Module):
      """Gate Unit
      g = sigmoid(Wx)
      x = g * x
      """
      def __init__(self, input_size):
          super(Gate, self).__init__()
          self.linear = nn.Linear(input_size, input_size, bias=False)
  
      def forward(self, x):
          """
          Args:
              x: batch * len * dim
              x_mask: batch * len (1 for padding, 0 for true)
          Output:
              res: batch * len * dim
          """
          x_proj = self.linear(x)
          gate = F.sigmoid(x)
          return x_proj * gate
  
  
  class SFU(nn.Module):
      """Semantic Fusion Unit
      The ouput vector is expected to not only retrieve correlative information from fusion vectors,
      but also retain partly unchange as the input vector
      """
      def __init__(self, input_size, fusion_size):
          super(SFU, self).__init__()
          self.linear_r = nn.Linear(input_size + fusion_size, input_size)
          self.linear_g = nn.Linear(input_size + fusion_size, input_size)
  
      def forward(self, x, fusions):
          r_f = torch.cat([x, fusions], 2)
          r = F.tanh(self.linear_r(r_f))
          g = F.sigmoid(self.linear_g(r_f))
          o = g * r + (1-g) * x
          return o
          
  
  # ------------------------------------------------------------------------------
  # Functional
  # ------------------------------------------------------------------------------
  
  
  def uniform_weights(x, x_mask):
      """Return uniform weights over non-masked x (a sequence of vectors).
  
      Args:
          x: batch * len * hdim
          x_mask: batch * len (1 for padding, 0 for true)
      Output:
          x_avg: batch * hdim
      """
      alpha = Variable(torch.ones(x.size(0), x.size(1)))
      if x.data.is_cuda:
          alpha = alpha.cuda()
      alpha = alpha * x_mask.eq(0).float()
      alpha = alpha / alpha.sum(1).expand(alpha.size())
      return alpha
  
  
  def weighted_avg(x, weights):
      """Return a weighted average of x (a sequence of vectors).
  
      Args:
          x: batch * len * hdim
          weights: batch * len, sum(dim = 1) = 1
      Output:
          x_avg: batch * hdim
      """
      return weights.unsqueeze(1).bmm(x).squeeze(1)
  