(self, x, y, mask)
return x
Shapes:
      mask: [nodes/edges, edges/nodes]
      q: [nodes/edges, h]
      k: [edges/nodes, h]
      v: [edges/nodes, h]
      score: [nodes/edges, edges/nodes]
      x_atten: [nodes/edges, h]
Shapes:
      mask: [nodes/edges, edges/nodes]
      q: [nodes/edges, h]
      k: [edges/nodes, h]
      v: [edges/nodes, h]
      score: [nodes/edges, edges/nodes]
      x_atten: [nodes/edges, h]
Shapes : mask : [ nodes / edges edges / nodes ] q : [ nodes / edges h ] k : [ edges / nodes h ] v : [ edges / nodes h ] score : [ nodes / edges edges / nodes ] x_atten : [ nodes / edges h ]
def forward(self, x, y, mask):
    """
    Shapes:
      mask: [nodes/edges, edges/nodes]
      q: [nodes/edges, h]
      k: [edges/nodes, h]
      v: [edges/nodes, h]
      score: [nodes/edges, edges/nodes]
      x_atten: [nodes/edges, h]
    """
    q = self.linear_q(x)
    k = self.linear_k(y)
    v = self.linear_v(y)
    score = torch.mm(q, k.t()) / math.sqrt(self.size)
    score = self.masked_softmax(score, mask, dim=1)
    x_atten = torch.mm(score, v)
    # dropout
    x_atten = self.dropout(x_atten)
    x = self.layer_norm_1(x + x_atten)
    x_linear = self.feed_forward(x)
    # dropout
    x_linear = self.dropout(x_linear)
    x = self.layer_norm_2(x + x_linear)
    return x
def forward ( self , x , y , mask ) : q = self . linear_q ( x ) k = self . linear_k ( y ) v = self . linear_v ( y ) score = torch . mm ( q , k . t ( ) ) / math . sqrt ( self . size ) score = self . masked_softmax ( score , mask , dim = 1 ) x_atten = torch . mm ( score , v ) # dropout x_atten = self . dropout ( x_atten ) x = self . layer_norm_1 ( x + x_atten ) x_linear = self . feed_forward ( x ) # dropout x_linear = self . dropout ( x_linear ) x = self . layer_norm_2 ( x + x_linear ) return x