# Path: nmt_chainer/models/feedforward/utils.py
# def apply_linear_layer_to_last_dims(Q, w_Q):
#     mb_size_Q, n_Q, d_model_Q = Q.data.shape
#     return F.reshape(w_Q(F.reshape(Q, (mb_size_Q * n_Q, d_model_Q))), (mb_size_Q, n_Q, -1))
# 
# class DropoutAndAddAndNormalize(Chain):
#     def __init__(self, dropout=None, residual_mode="normal", no_normalize=False):
#         super(DropoutAndAddAndNormalize, self).__init__()
#         
#         if not no_normalize:
#             LayerNormalization = get_layer_normalization_class()
#             self.add_link("normalizing_layer", LayerNormalization())
#         
#         assert residual_mode in "normal none after".split()
#         self.residual_mode = residual_mode
#         self.no_normalize = no_normalize
#         self.dropout = dropout
#         
#         
#     def apply_layer_normalization(self, added_output):
#         if len(added_output.shape) > 2:
#             d_model = added_output.shape[-1]
#             final_layer = F.reshape(
#                 self.normalizing_layer(
#                     F.reshape(added_output, (-1, d_model))
#                     ), added_output.shape)                
#         else:
#             final_layer = self.normalizing_layer(added_output)
#             
#         return final_layer
#             
#     def __call__(self, sub_output, inpt):
#         if self.dropout is not None:
#             sub_output = F.dropout(sub_output, ratio=self.dropout)
#             
#         if self.residual_mode == "normal":
#             added_output = sub_output + inpt
#         else:
#             added_output = sub_output
#         
#         if self.no_normalize:
#             final_layer = added_output
#         else:
#             final_layer = self.apply_layer_normalization(added_output)
# 
#         if self.residual_mode == "after":
#             final_layer = final_layer + inpt
# 
#         return final_layer


import chainer
import numpy as np
import chainer.functions as F
import chainer.links as L
from chainer import Variable, Chain, ChainList
from nmt_chainer.models.feedforward.utils import apply_linear_layer_to_last_dims, DropoutAndAddAndNormalize
        if self.experimental_relu:
            addressing_weights = F.relu(scaled_scalar_product)
        else:
            addressing_weights = F.reshape(F.softmax(F.reshape(scaled_scalar_product, (mb_size * n_Q * self.n_heads, seq_length_K)),
                                                     #use_cudnn=disable_cudnn_softmax
                                                     ),
                                           (mb_size, self.n_heads, n_Q, seq_length_K) )
        
        if self.dropout is not None:
            addressing_weights = F.dropout(addressing_weights, ratio=self.dropout)
        
#         print "A", addressing_weights.data
        reorganized_V = reorganize_by_head(proj_V, self.n_heads)
        reorganized_result = batch_matmul_last_dims(addressing_weights, reorganized_V)
        result = undo_reorganize_by_head(reorganized_result)
        
        if self.n_heads >= 2:
            result = apply_linear_layer_to_last_dims(result, self.w_O)
        
#         print "R", result.data
        return result
    
    
class AddAndNormalizedAttentionBase(Chain):
    def __init__(self, d_model, n_heads, experimental_relu=False, dropout=None, residual_mode="normal", no_normalize=False):
        super(AddAndNormalizedAttentionBase, self).__init__(
            multi_attention= ConstantSizeMultiBatchMultiHeadAttention(d_model = d_model, n_heads=n_heads,
                                                             experimental_relu=experimental_relu,
                                                                     dropout=dropout),
                                                            

# Path: nmt_chainer/models/feedforward/utils.py
# def apply_linear_layer_to_last_dims(Q, w_Q):
#     mb_size_Q, n_Q, d_model_Q = Q.data.shape
#     return F.reshape(w_Q(F.reshape(Q, (mb_size_Q * n_Q, d_model_Q))), (mb_size_Q, n_Q, -1))
# 
# class DropoutAndAddAndNormalize(Chain):
#     def __init__(self, dropout=None, residual_mode="normal", no_normalize=False):
#         super(DropoutAndAddAndNormalize, self).__init__()
#         
#         if not no_normalize:
#             LayerNormalization = get_layer_normalization_class()
#             self.add_link("normalizing_layer", LayerNormalization())
#         
#         assert residual_mode in "normal none after".split()
#         self.residual_mode = residual_mode
#         self.no_normalize = no_normalize
#         self.dropout = dropout
#         
#         
#     def apply_layer_normalization(self, added_output):
#         if len(added_output.shape) > 2:
#             d_model = added_output.shape[-1]
#             final_layer = F.reshape(
#                 self.normalizing_layer(
#                     F.reshape(added_output, (-1, d_model))
#                     ), added_output.shape)                
#         else:
#             final_layer = self.normalizing_layer(added_output)
#             
#         return final_layer
#             
#     def __call__(self, sub_output, inpt):
#         if self.dropout is not None:
#             sub_output = F.dropout(sub_output, ratio=self.dropout)
#             
#         if self.residual_mode == "normal":
#             added_output = sub_output + inpt
#         else:
#             added_output = sub_output
#         
#         if self.no_normalize:
#             final_layer = added_output
#         else:
#             final_layer = self.apply_layer_normalization(added_output)
# 
#         if self.residual_mode == "after":
#             final_layer = final_layer + inpt
# 
#         return final_layer


# Path: nmt_chainer/models/feedforward/multi_attention.py
import chainer
import numpy as np
import chainer.functions as F
import chainer.links as L
from chainer import Variable, Chain, ChainList
from nmt_chainer.models.feedforward.utils import apply_linear_layer_to_last_dims, DropoutAndAddAndNormalize

        if self.experimental_relu:
            addressing_weights = F.relu(scaled_scalar_product)
        else:
            addressing_weights = F.reshape(F.softmax(F.reshape(scaled_scalar_product, (mb_size * n_Q * self.n_heads, seq_length_K)),
                                                     #use_cudnn=disable_cudnn_softmax
                                                     ),
                                           (mb_size, self.n_heads, n_Q, seq_length_K) )
        
        if self.dropout is not None:
            addressing_weights = F.dropout(addressing_weights, ratio=self.dropout)
        
#         print "A", addressing_weights.data
        reorganized_V = reorganize_by_head(proj_V, self.n_heads)
        reorganized_result = batch_matmul_last_dims(addressing_weights, reorganized_V)
        result = undo_reorganize_by_head(reorganized_result)
        
        if self.n_heads >= 2:
            result = apply_linear_layer_to_last_dims(result, self.w_O)
        
#         print "R", result.data
        return result
    
    
class AddAndNormalizedAttentionBase(Chain):
    def __init__(self, d_model, n_heads, experimental_relu=False, dropout=None, residual_mode="normal", no_normalize=False):
        super(AddAndNormalizedAttentionBase, self).__init__(
            multi_attention= ConstantSizeMultiBatchMultiHeadAttention(d_model = d_model, n_heads=n_heads,
                                                             experimental_relu=experimental_relu,
                                                                     dropout=dropout),
                                                            

            residual_layer = DropoutAndAddAndNormalize(dropout=dropout, residual_mode=residual_mode, no_normalize=no_normalize)