# Path: nmt_chainer/models/feedforward/utils.py
# def generate_pos_vectors(d_model, max_length):
#     pos_component = np.arange(max_length, dtype = np.float32)
#     dim_component = np.arange(d_model, dtype = np.float32)
#     dim_component_even = np.floor_divide(dim_component, 2) * 2
#     dim_factor = np.power(1e-4, dim_component_even / d_model)
#     pos_dim = pos_component[:, None] * dim_factor[None, :]
#     pos_dim[:, ::2] = np.sin(pos_dim[:, ::2])
#     pos_dim[:, 1::2] = np.cos(pos_dim[:, 1::2])
#     return pos_dim
# 
# def make_batch_mask(mb_size, n_head, max_length_1, max_length_2, 
#                     key_seq_lengths=None,
#                     future_mask=False,
#                     mask_value=-10000):
#     
#     if future_mask:
#         assert max_length_1 == max_length_2
#         mask = np.array(
#                 np.broadcast_to(( (-mask_value) * (np.tri(max_length_1, dtype = np.float32)-1))[None,None,:,:], 
#                                 (mb_size, n_head, max_length_1, max_length_2))
#                 )
#     else:
#         mask = np.zeros((mb_size, n_head, max_length_1, max_length_2), dtype = np.float32)
#         
#     if key_seq_lengths is not None:
#         assert mb_size == len(key_seq_lengths)
#         assert min(key_seq_lengths) > 0
#         assert max(key_seq_lengths) <= max_length_2
#         for num_batch, length in enumerate(key_seq_lengths):
#             mask[num_batch, :, :, length:] = mask_value
# 
#     return mask
# 
# def pad_data(data, pad_value=0, add_eos=None):
#     mb_size = len(data)
#     padded_length = max(len(x) for x in data)
#     if add_eos is not None:
#         padded_length += 1
#     padded_array = np.zeros((mb_size, padded_length), dtype=np.int32)
#     if pad_value != 0:
#         padded_array += pad_value
#     for num_batch, seq in enumerate(data):
#         padded_array[num_batch, :len(seq)] = seq
#         if add_eos is not None:
#             padded_array[num_batch, len(seq)] = add_eos
#     return padded_array
# 
# class FeedForward(Chain):
#     def __init__(self, d_model=512, d_ff=2048, dropout=None,
#             residual_mode="normal", no_normalize=False):
#         super(FeedForward, self).__init__(
#             lin1 = L.Linear(d_model, d_ff),
#             lin2 = L.Linear(d_ff, d_model),
#             layer_reduce = DropoutAndAddAndNormalize(dropout=dropout, residual_mode=residual_mode, no_normalize=no_normalize)
#         )
#         
# #         self.dropout = dropout
# #         
# #         if not no_normalize:
# #             self.add_link("normalization_layer", LayerNormalization())
# #         
# #         self.no_add = no_add
# #         self.no_normalize = no_normalize
#         
#     def __call__(self, x_input):
# #         print "FF", x_input.data
#         if len(x_input.data.shape) > 2:
#             x = F.reshape(x_input, (-1, x_input.shape[-1]))
#         else:
#             x = x_input
#             
#         ff_output = self.lin2(F.relu(self.lin1(x)))
#         
#         norm_ff_output = self.layer_reduce(ff_output, x)
#         
# #         if self.dropout is not None:
# #             ff_output = F.dropout(ff_output, self.dropout, train=train)
# #             
# #         if self.no_add:
# #             added_output = ff_output
# #         else:
# #             added_output = ff_output + x
# #         
# #         if self.no_normalize:
# #             norm_ff_output = added_output
# #         else:
# #             norm_ff_output = self.normalization_layer(added_output)
#         
#         if len(x_input.data.shape) > 2:
#             norm_ff_output = F.reshape(norm_ff_output, x_input.data.shape)
#             
# #         print "FFR", norm_ff_output.data
#         return norm_ff_output
# 
# Path: nmt_chainer/models/feedforward/multi_attention.py
# class AddAndNormalizedSelfAttentionLayer(AddAndNormalizedAttentionBase):
#     def __init__(self, d_model, n_heads, experimental_relu=False, dropout=None, residual_mode="normal", no_normalize=False):
#         super(AddAndNormalizedSelfAttentionLayer, self).__init__(
#             d_model, n_heads, experimental_relu=experimental_relu, dropout=dropout,
#             residual_mode=residual_mode, no_normalize=no_normalize
#         )
#         
#     def __call__(self, x, mask, only_last=False):
# #         print "SELF"
#         if only_last:
#             x_in = self.extract_last(x)
#         else:
#             x_in = x
#         sub_output = self.multi_attention(x_in, x, x, mask)
#             
#         return self.residual_layer(sub_output, x_in)


import chainer
import numpy as np
import chainer.functions as F
import chainer.links as L
import six
from chainer import Variable, Chain, ChainList
from nmt_chainer.models.feedforward.utils import generate_pos_vectors, make_batch_mask, pad_data, FeedForward
from nmt_chainer.models.feedforward.multi_attention import AddAndNormalizedSelfAttentionLayer
#!/usr/bin/env python
from __future__ import absolute_import, division, print_function, unicode_literals



class EncoderLayer(Chain):
    def __init__(self, d_model, n_heads, d_ff=2048, experimental_relu=False, dropout=None, residual_mode="normal", no_normalize=False):
        super(EncoderLayer, self).__init__(
            ff_layer = FeedForward(d_model, d_ff=d_ff, dropout=dropout, residual_mode=residual_mode, no_normalize=no_normalize),

# Path: nmt_chainer/models/feedforward/utils.py
# def generate_pos_vectors(d_model, max_length):
#     pos_component = np.arange(max_length, dtype = np.float32)
#     dim_component = np.arange(d_model, dtype = np.float32)
#     dim_component_even = np.floor_divide(dim_component, 2) * 2
#     dim_factor = np.power(1e-4, dim_component_even / d_model)
#     pos_dim = pos_component[:, None] * dim_factor[None, :]
#     pos_dim[:, ::2] = np.sin(pos_dim[:, ::2])
#     pos_dim[:, 1::2] = np.cos(pos_dim[:, 1::2])
#     return pos_dim
# 
# def make_batch_mask(mb_size, n_head, max_length_1, max_length_2, 
#                     key_seq_lengths=None,
#                     future_mask=False,
#                     mask_value=-10000):
#     
#     if future_mask:
#         assert max_length_1 == max_length_2
#         mask = np.array(
#                 np.broadcast_to(( (-mask_value) * (np.tri(max_length_1, dtype = np.float32)-1))[None,None,:,:], 
#                                 (mb_size, n_head, max_length_1, max_length_2))
#                 )
#     else:
#         mask = np.zeros((mb_size, n_head, max_length_1, max_length_2), dtype = np.float32)
#         
#     if key_seq_lengths is not None:
#         assert mb_size == len(key_seq_lengths)
#         assert min(key_seq_lengths) > 0
#         assert max(key_seq_lengths) <= max_length_2
#         for num_batch, length in enumerate(key_seq_lengths):
#             mask[num_batch, :, :, length:] = mask_value
# 
#     return mask
# 
# def pad_data(data, pad_value=0, add_eos=None):
#     mb_size = len(data)
#     padded_length = max(len(x) for x in data)
#     if add_eos is not None:
#         padded_length += 1
#     padded_array = np.zeros((mb_size, padded_length), dtype=np.int32)
#     if pad_value != 0:
#         padded_array += pad_value
#     for num_batch, seq in enumerate(data):
#         padded_array[num_batch, :len(seq)] = seq
#         if add_eos is not None:
#             padded_array[num_batch, len(seq)] = add_eos
#     return padded_array
# 
# class FeedForward(Chain):
#     def __init__(self, d_model=512, d_ff=2048, dropout=None,
#             residual_mode="normal", no_normalize=False):
#         super(FeedForward, self).__init__(
#             lin1 = L.Linear(d_model, d_ff),
#             lin2 = L.Linear(d_ff, d_model),
#             layer_reduce = DropoutAndAddAndNormalize(dropout=dropout, residual_mode=residual_mode, no_normalize=no_normalize)
#         )
#         
# #         self.dropout = dropout
# #         
# #         if not no_normalize:
# #             self.add_link("normalization_layer", LayerNormalization())
# #         
# #         self.no_add = no_add
# #         self.no_normalize = no_normalize
#         
#     def __call__(self, x_input):
# #         print "FF", x_input.data
#         if len(x_input.data.shape) > 2:
#             x = F.reshape(x_input, (-1, x_input.shape[-1]))
#         else:
#             x = x_input
#             
#         ff_output = self.lin2(F.relu(self.lin1(x)))
#         
#         norm_ff_output = self.layer_reduce(ff_output, x)
#         
# #         if self.dropout is not None:
# #             ff_output = F.dropout(ff_output, self.dropout, train=train)
# #             
# #         if self.no_add:
# #             added_output = ff_output
# #         else:
# #             added_output = ff_output + x
# #         
# #         if self.no_normalize:
# #             norm_ff_output = added_output
# #         else:
# #             norm_ff_output = self.normalization_layer(added_output)
#         
#         if len(x_input.data.shape) > 2:
#             norm_ff_output = F.reshape(norm_ff_output, x_input.data.shape)
#             
# #         print "FFR", norm_ff_output.data
#         return norm_ff_output
# 
# Path: nmt_chainer/models/feedforward/multi_attention.py
# class AddAndNormalizedSelfAttentionLayer(AddAndNormalizedAttentionBase):
#     def __init__(self, d_model, n_heads, experimental_relu=False, dropout=None, residual_mode="normal", no_normalize=False):
#         super(AddAndNormalizedSelfAttentionLayer, self).__init__(
#             d_model, n_heads, experimental_relu=experimental_relu, dropout=dropout,
#             residual_mode=residual_mode, no_normalize=no_normalize
#         )
#         
#     def __call__(self, x, mask, only_last=False):
# #         print "SELF"
#         if only_last:
#             x_in = self.extract_last(x)
#         else:
#             x_in = x
#         sub_output = self.multi_attention(x_in, x, x, mask)
#             
#         return self.residual_layer(sub_output, x_in)


# Path: nmt_chainer/models/feedforward/encoder.py
import chainer
import numpy as np
import chainer.functions as F
import chainer.links as L
import six
from chainer import Variable, Chain, ChainList
from nmt_chainer.models.feedforward.utils import generate_pos_vectors, make_batch_mask, pad_data, FeedForward
from nmt_chainer.models.feedforward.multi_attention import AddAndNormalizedSelfAttentionLayer

#!/usr/bin/env python
from __future__ import absolute_import, division, print_function, unicode_literals



class EncoderLayer(Chain):
    def __init__(self, d_model, n_heads, d_ff=2048, experimental_relu=False, dropout=None, residual_mode="normal", no_normalize=False):
        super(EncoderLayer, self).__init__(
            ff_layer = FeedForward(d_model, d_ff=d_ff, dropout=dropout, residual_mode=residual_mode, no_normalize=no_normalize),

            self_attention_layer = AddAndNormalizedSelfAttentionLayer(d_model=d_model, n_heads=n_heads,