(self, query, key, value, mask, head_mask = None)
Parameters
        ----------
        query: torch.tensor(bs, seq_length, dim)
        key: torch.tensor(bs, seq_length, dim)
        value: torch.tensor(bs, seq_length, dim)
        mask: torch.tensor(bs, seq_length)

        Outputs
        -------
        weights: torch.tensor(bs, n_heads, seq_length, seq_length)
            Attention weights
        context: torch.tensor(bs, seq_length, dim)
            Contextualized layer. Optional: only if `output_attentions=True`
Parameters
        ----------
        query: torch.tensor(bs, seq_length, dim)
        key: torch.tensor(bs, seq_length, dim)
        value: torch.tensor(bs, seq_length, dim)
        mask: torch.tensor(bs, seq_length)
Parameters ---------- query : torch . tensor ( bs seq_length dim ) key : torch . tensor ( bs seq_length dim ) value : torch . tensor ( bs seq_length dim ) mask : torch . tensor ( bs seq_length )
def forward(self, query, key, value, mask, head_mask = None):
        """
        Parameters
        ----------
        query: torch.tensor(bs, seq_length, dim)
        key: torch.tensor(bs, seq_length, dim)
        value: torch.tensor(bs, seq_length, dim)
        mask: torch.tensor(bs, seq_length)

        Outputs
        -------
        weights: torch.tensor(bs, n_heads, seq_length, seq_length)
            Attention weights
        context: torch.tensor(bs, seq_length, dim)
            Contextualized layer. Optional: only if `output_attentions=True`
        """
        bs, q_length, dim = query.size()
        k_length = key.size(1)
        # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
        # assert key.size() == value.size()

        dim_per_head = self.dim // self.n_heads

        mask_reshp = (bs, 1, 1, k_length)

        def shape(x):
            """ separate heads """
            return x.view(bs, -1, self.n_heads, dim_per_head).transpose(1, 2)

        def unshape(x):
            """ group heads """
            return x.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * dim_per_head)

        q = shape(self.q_lin(query))           # (bs, n_heads, q_length, dim_per_head)
        k = shape(self.k_lin(key))             # (bs, n_heads, k_length, dim_per_head)
        v = shape(self.v_lin(value))           # (bs, n_heads, k_length, dim_per_head)

        q = q / math.sqrt(dim_per_head)                     # (bs, n_heads, q_length, dim_per_head)
        scores = torch.matmul(q, k.transpose(2,3))          # (bs, n_heads, q_length, k_length)
        mask = (mask==0).view(mask_reshp).expand_as(scores) # (bs, n_heads, q_length, k_length)
        scores.masked_fill_(mask, -float('inf'))            # (bs, n_heads, q_length, k_length)

        weights = nn.Softmax(dim=-1)(scores)   # (bs, n_heads, q_length, k_length)
        weights = self.dropout(weights)        # (bs, n_heads, q_length, k_length)

        # Mask heads if we want to
        if head_mask is not None:
            weights = weights * head_mask

        context = torch.matmul(weights, v)     # (bs, n_heads, q_length, dim_per_head)
        context = unshape(context)             # (bs, q_length, dim)
        context = self.out_lin(context)        # (bs, q_length, dim)

        if self.output_attentions:
            return (context, weights)
        else:
            return (context,)
def forward ( self , query , key , value , mask , head_mask = None ) : bs , q_length , dim = query . size ( ) k_length = key . size ( 1 ) # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim) # assert key.size() == value.size() dim_per_head = self . dim // self . n_heads mask_reshp = ( bs , 1 , 1 , k_length ) def shape ( x ) : """ separate heads """ return x . view ( bs , - 1 , self . n_heads , dim_per_head ) . transpose ( 1 , 2 ) def unshape ( x ) : """ group heads """ return x . transpose ( 1 , 2 ) . contiguous ( ) . view ( bs , - 1 , self . n_heads * dim_per_head ) q = shape ( self . q_lin ( query ) ) # (bs, n_heads, q_length, dim_per_head) k = shape ( self . k_lin ( key ) ) # (bs, n_heads, k_length, dim_per_head) v = shape ( self . v_lin ( value ) ) # (bs, n_heads, k_length, dim_per_head) q = q / math . sqrt ( dim_per_head ) # (bs, n_heads, q_length, dim_per_head) scores = torch . matmul ( q , k . transpose ( 2 , 3 ) ) # (bs, n_heads, q_length, k_length) mask = ( mask == 0 ) . view ( mask_reshp ) . expand_as ( scores ) # (bs, n_heads, q_length, k_length) scores . masked_fill_ ( mask , - float ( 'inf' ) ) # (bs, n_heads, q_length, k_length) weights = nn . Softmax ( dim = - 1 ) ( scores ) # (bs, n_heads, q_length, k_length) weights = self . dropout ( weights ) # (bs, n_heads, q_length, k_length) # Mask heads if we want to if head_mask is not None : weights = weights * head_mask context = torch . matmul ( weights , v ) # (bs, n_heads, q_length, dim_per_head) context = unshape ( context ) # (bs, q_length, dim) context = self . out_lin ( context ) # (bs, q_length, dim) if self . output_attentions : return ( context , weights ) else : return ( context , )