text: """
  decoder.py
  
  Transformer's decoder.
  """
  from transformer.shared import *
  
  
  class Decoder(nn.Module):
      def __init__(self, output_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, device, max_length=70):
          super().__init__()
          
          self.device = device
          
          self.tok_embedding = nn.Embedding(output_dim, hid_dim)
          self.pos_embedding = nn.Embedding(max_length, hid_dim)
          
          self.layers = nn.ModuleList([DecoderLayer(hid_dim,
                                                    n_heads,
                                                    pf_dim,
                                                    dropout,
                                                    device)
                                       for _ in range(n_layers)])
          
          self.fc_out = nn.Linear(hid_dim, output_dim)
          
          self.dropout = nn.Dropout(dropout)
          
          self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)
      
      def forward(self, trg, enc_src, trg_mask, src_mask):
          # trg = [batch size, trg len]
          # enc_src = [batch size, src len, hid dim]
          # trg_mask = [batch size, trg len]
          # src_mask = [batch size, src len]
          
          batch_size = trg.shape[0]
          trg_len = trg.shape[1]
          
          pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)
          
          # pos = [batch size, trg len]
          
          trg = self.dropout((self.tok_embedding(trg) * self.scale) + self.pos_embedding(pos))
          
          # trg = [batch size, trg len, hid dim]
          
          for layer in self.layers:
              trg, attention = layer(trg, enc_src, trg_mask, src_mask)
          
          # trg = [batch size, trg len, hid dim]
          # attention = [batch size, n heads, trg len, src len]
          
          output = self.fc_out(trg)
          
          # output = [batch size, trg len, output dim]
          
          return output, attention
  
  
  class DecoderLayer(nn.Module):
      def __init__(self,
                   hid_dim,
                   n_heads,
                   pf_dim,
                   dropout,
                   device):
          super().__init__()
          
          self.layer_norm = nn.LayerNorm(hid_dim)
          self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)
          self.encoder_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)
          self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim,
                                                                       pf_dim,
                                                                       dropout)
          self.dropout = nn.Dropout(dropout)
      
      def forward(self, trg, enc_src, trg_mask, src_mask):
          # trg = [batch size, trg len, hid dim]
          # enc_src = [batch size, src len, hid dim]
          # trg_mask = [batch size, trg len]
          # src_mask = [batch size, src len]
          
          # self attention
          _trg, _ = self.self_attention(trg, trg, trg, trg_mask)
          
          # dropout, residual connection and layer norm
          trg = self.layer_norm(trg + self.dropout(_trg))
          
          # trg = [batch size, trg len, hid dim]
          
          # encoder attention
          _trg, attention = self.encoder_attention(trg, enc_src, enc_src, src_mask)
          
          # dropout, residual connection and layer norm
          trg = self.layer_norm(trg + self.dropout(_trg))
          
          # trg = [batch size, trg len, hid dim]
          
          # positionwise feedforward
          _trg = self.positionwise_feedforward(trg)
          
          # dropout, residual and layer norm
          trg = self.layer_norm(trg + self.dropout(_trg))
          
          # trg = [batch size, trg len, hid dim]
          # attention = [batch size, n heads, trg len, src len]
          
          return trg, attention
  