import torch
import torch.nn as nn
import torch.nn.init as init
import numpy as np

__author__ = "Yu-Hsiang Huang"

class Linear(nn.Module):
    ''' Simple Linear layer with xavier init '''
    def __init__(self, d_in, d_out, bias=True):
        super(Linear, self).__init__()
        self.linear = nn.Linear(d_in, d_out, bias=bias)
        init.xavier_normal(self.linear.weight)

    def forward(self, x):
        return self.linear(x)

class Bottle(nn.Module):
    ''' Perform the reshape routine before and after an operation '''

    def forward(self, input):
        if len(input.size()) <= 2:
            return super(Bottle, self).forward(input)
        size = input.size()[:2]
        out = super(Bottle, self).forward(input.view(size[0]*size[1], -1))
        return out.view(size[0], size[1], -1)

class BottleLinear(Bottle, Linear):
    ''' Perform the reshape routine before and after a linear projection '''
    pass

class LayerNormalization(nn.Module):
    ''' Layer normalization module '''

    def __init__(self, d_hid, eps=1e-3):
        super(LayerNormalization, self).__init__()

        self.eps = eps
        self.a_2 = nn.Parameter(torch.ones(d_hid), requires_grad=True)
        self.b_2 = nn.Parameter(torch.zeros(d_hid), requires_grad=True)

    def forward(self, z):
        if z.size(1) == 1:
            return z

        mu = torch.mean(z, dim=1)
        sigma = torch.std(z, dim=1)
        ln_out = (z - mu.expand_as(z)) / (sigma.expand_as(z) + self.eps)
        ln_out = ln_out * self.a_2.expand_as(ln_out) + self.b_2.expand_as(ln_out)

        return ln_out

class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, d_model, attn_dropout=0.1):
        super(ScaledDotProductAttention, self).__init__()
        self.temper = np.power(d_model, 0.5)
        self.dropout = nn.Dropout(attn_dropout)
        self.softmax = nn.Softmax()

    def forward(self, q, k, v, attn_mask=None):
        _, len_q, _ = q.size()
        _, len_k, _ = k.size()

        attn = torch.bmm(q, k.transpose(1, 2)) / self.temper

        if attn_mask is not None:
            attn.data.masked_fill_(attn_mask, -float('inf'))

        attn = self.softmax(attn.view(-1, len_k)).view(-1, len_q, len_k)
        attn = self.dropout(attn)
        output = torch.bmm(attn, v)
        return output, attn

@@ -923,16 +923,143 @@
   pass%0A%0A
+class BottleSoftmax(Bottle, nn.Softmax):%0A    ''' Perform the reshape routine before and after a softmax operation'''%0A    pass%0A%0A
 class La
@@ -2022,19 +2022,22 @@
 ftmax = 
-nn.
+Bottle
 Softmax(
@@ -2090,70 +2090,8 @@
 ne):
-%0A        _, len_q, _ = q.size()%0A        _, len_k, _ = k.size()
 %0A%0A  
@@ -2184,16 +2184,253 @@
 t None:%0A
+%0A            assert attn_mask.size() == attn.size(), %5C%0A                    'Attention mask shape %7B%7D mismatch ' %5C%0A                    'with Attention logit tensor shape ' %5C%0A                    '%7B%7D.'.format(attn_mask.size(), attn.size())%0A%0A
         
@@ -2519,47 +2519,8 @@
 attn
-.view(-1, len_k)).view(-1, len_q, len_k
 )%0A  
@@ -2587,16 +2587,17 @@
 ttn, v)%0A
+%0A
         

add attn mask shape check