(self, x)
| 38 | self.mask = a_dense |
| 39 | |
| 40 | def forward(self, x): |
| 41 | b, n, s = x.shape |
| 42 | x = x.reshape(b*n, s) |
| 43 | |
| 44 | atts_stack = [] |
| 45 | # multi-heads attention |
| 46 | for n in range(self.head_n): |
| 47 | h = torch.matmul(x, self.heads_dict[n, 0]) |
| 48 | edge_h = torch.cat((h[self.edges[0, :], :], h[self.edges[1, :], :]), dim=1).t() # [Ni, Nj] |
| 49 | atts = self.heads_dict[n, 1].mm(edge_h).squeeze() |
| 50 | atts = self.leakyrelu(atts) |
| 51 | atts_stack.append(atts) |
| 52 | |
| 53 | mt_atts = torch.stack(atts_stack, dim=1) |
| 54 | mt_atts = self.linear(mt_atts) |
| 55 | new_values = self.values * mt_atts.squeeze() |
| 56 | atts_mat = torch.sparse_coo_tensor(self.edges, new_values) |
| 57 | atts_mat = atts_mat.to_dense() + self.mask |
| 58 | atts_mat = self.softmax(atts_mat) |
| 59 | return atts_mat |
| 60 | |
| 61 | |
| 62 | class MLP(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected