# 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
        self.dropout = dropout
                                                 
                                                 
    def __call__(self, Q, K, V, batch_mask = None):
#         print "Q",
#         print Q.data
#         print "K",
#         print K.data
#         print "V",
#         print V.data
#         print "M", batch_mask
        
        mb_size_Q, n_Q, d_model_Q = Q.data.shape
        mb_size_K, seq_length_K, d_model_K = K.data.shape
        mb_size_V, seq_length_V, d_model_V = V.data.shape
                                            
        assert mb_size_Q == mb_size_K == mb_size_V
        assert d_model_Q == d_model_K == d_model_V == self.d_model
        assert seq_length_K == seq_length_V                
                                            
        mb_size = mb_size_Q
                                            
        if batch_mask is not None:                            
            mb_size_batch_mask, mask_n_heads, mask_n_Q, mask_seq_length_K = batch_mask.shape
            assert  mb_size_batch_mask == mb_size
            assert mask_n_heads == self.n_heads
            assert mask_n_Q == n_Q, "%i != %i"%(mask_n_Q, n_Q)
            assert mask_seq_length_K == mask_seq_length_K               
                                            
        

# 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

        self.dropout = dropout
                                                 
                                                 
    def __call__(self, Q, K, V, batch_mask = None):
#         print "Q",
#         print Q.data
#         print "K",
#         print K.data
#         print "V",
#         print V.data
#         print "M", batch_mask
        
        mb_size_Q, n_Q, d_model_Q = Q.data.shape
        mb_size_K, seq_length_K, d_model_K = K.data.shape
        mb_size_V, seq_length_V, d_model_V = V.data.shape
                                            
        assert mb_size_Q == mb_size_K == mb_size_V
        assert d_model_Q == d_model_K == d_model_V == self.d_model
        assert seq_length_K == seq_length_V                
                                            
        mb_size = mb_size_Q
                                            
        if batch_mask is not None:                            
            mb_size_batch_mask, mask_n_heads, mask_n_Q, mask_seq_length_K = batch_mask.shape
            assert  mb_size_batch_mask == mb_size
            assert mask_n_heads == self.n_heads
            assert mask_n_Q == n_Q, "%i != %i"%(mask_n_Q, n_Q)
            assert mask_seq_length_K == mask_seq_length_K               
                                            
        

        proj_Q = apply_linear_layer_to_last_dims(Q, self.w_Q)