Input: x: default node feature. depth: The depth of the node in the AST. Output: emb_dim-dimensional vector
| 3 | |
| 4 | |
| 5 | class ASTNodeEncoder(torch.nn.Module): |
| 6 | ''' |
| 7 | Input: |
| 8 | x: default node feature. |
| 9 | depth: The depth of the node in the AST. |
| 10 | |
| 11 | Output: |
| 12 | emb_dim-dimensional vector |
| 13 | |
| 14 | ''' |
| 15 | def __init__(self, emb_dim, max_depth, enc_dims=[]): |
| 16 | super(ASTNodeEncoder, self).__init__() |
| 17 | |
| 18 | self.max_depth = max_depth |
| 19 | if enc_dims: |
| 20 | edim = emb_dim//(len(enc_dims) + 1) |
| 21 | l = [torch.nn.Embedding(n, edim) for n in enc_dims] |
| 22 | self.embs = torch.nn.ModuleList(l) |
| 23 | rest = emb_dim - (edim * (len(enc_dims) + 1)) |
| 24 | self.depth_encoder = torch.nn.Embedding(self.max_depth + 1, edim+rest) |
| 25 | else: |
| 26 | self.embs = None |
| 27 | self.depth_encoder = torch.nn.Embedding(self.max_depth + 1, emb_dim) |
| 28 | |
| 29 | |
| 30 | def forward(self, x, depth): |
| 31 | depth[depth > self.max_depth] = self.max_depth |
| 32 | if self.embs is None: |
| 33 | return torch.cat([x, self.depth_encoder(depth)], dim=-1) |
| 34 | return torch.cat([enc(x[:, i]) for i, enc in enumerate(self.embs)]+[self.depth_encoder(depth)], dim=-1) |
| 35 | |
| 36 | |
| 37 | # VT: with ourgraphs the below "AST" edges etc do not have to be AST edges but may be arbitray (data flow etc) |