Generative Code Modeling with Graphs

1 INTRODUCTION . Learning to understand and generate programs is an important building block for procedural artificial intelligence and more intelligent software engineering tools . It is also an interesting task in the research of structured prediction methods : while imbued with formal semantics and strict syntactic rules , natural source code carries aspects of natural languages , since it acts as a means of communicating intent among developers . Early works in the area have shown that approaches from natural language processing can be applied successfully to source code ( Hindle et al. , 2012 ) , whereas the programming languages community has had successes in focusing exclusively on formal semantics . More recently , methods handling both modalities ( i.e. , the formal and natural language aspects ) have shown successes on important software engineering tasks ( Raychev et al. , 2015 ; Bichsel et al. , 2016 ; Allamanis et al. , 2018b ) and semantic parsing ( Yin & Neubig , 2017 ; Rabinovich et al. , 2017 ) . However , current generative models of source code mostly focus on only one of these modalities at a time . For example , program synthesis tools based on enumeration and deduction ( Solar-Lezama , 2008 ; Polozov & Gulwani , 2015 ; Feser et al. , 2015 ; Feng et al. , 2018 ) are successful at generating programs that satisfy some ( usually incomplete ) formal specification but are often obviously wrong on manual inspection , as they can not distinguish unlikely from likely , “ natural ” programs . On the other hand , learned code models have succeeded in generating realistic-looking programs ( Maddison & Tarlow , 2014 ; Bielik et al. , 2016 ; Parisotto et al. , 2017 ; Rabinovich et al. , 2017 ; Yin & Neubig , 2017 ) . However , these programs often fail to be semantically relevant , for example because variables are not used consistently . In this work , we try to overcome these challenges for generative code models and present a general method for generative models that can incorporate structured information that is deterministically available at generation time . We focus our attention on generating source code and follow the ideas of program graphs ( Allamanis et al. , 2018b ) that have been shown to learn semantically meaningful representations of ( pre-existing ) programs . To achieve this , we lift grammar-based tree decoder models into the graph setting , where the diverse relationships between various elements of the generated code can be modeled . For this , the syntax tree under generation is augmented with additional edges denoting known relationships ( e.g. , last use of variables ) . We then interleave the steps of the generative procedure with neural message passing ( Gilmer et al. , 2017 ) to compute more precise representations of the intermediate states of the program generation . This is fundamentally different from sequential generative models of graphs ( Li et al. , 2018 ; Samanta et al. , 2018 ) , which aim to generate all edges and nodes , whereas our graphs are deterministic augmentations of generated trees . To summarize , we present a ) a general graph-based generative procedure for highly structured objects , incorporating rich structural information ; b ) ExprGen , a new code generation task focused on Algorithm 1 Pseudocode for Expand Input : Context c , partial AST a , node v to expand 1 : hv ← getRepresentation ( c , a , v ) 2 : rhs ← pickProduction ( v , hv ) 3 : for child node type ` ∈ rhs do 4 : ( a , u ) ← insertChild ( a , ` ) 5 : if ` is nonterminal type then 6 : a← Expand ( c , a , u ) 7 : return a int ilOffsetIdx = Array.IndexOf ( sortedILOffsets , map.ILOffset ) ; int nextILOffsetIdx = ilOffsetIdx + 1 ; int nextMapILOffset = nextILOffsetIdx < sortedILOffsets.Length ? sortedILOffsets [ nextILOffsetIdx ] : int.MaxValue ; Figure 1 : Example for ExprGen , target ex- pression to be generated is marked . Taken from BenchmarkDotNet , lightly edited for formatting . generating small , but semantically complex expressions conditioned on source code context ; and c ) a comprehensive experimental evaluation of our generative procedure and a range of baseline methods from the literature . 2 BACKGROUND & TASK . The most general form of the code generation task is to produce a ( partial ) program in a programming language given some context information c. This context information can be natural language ( as in , e.g. , semantic parsing ) , input-output examples ( e.g. , inductive program synthesis ) , partial program sketches , etc . Early methods generate source code as a sequence of tokens ( Hindle et al. , 2012 ; Hellendoorn & Devanbu , 2017 ) and sometimes fail to produce syntactically correct code . More recent models are sidestepping this issue by using the target language ’ s grammar to generate abstract syntax trees ( ASTs ) ( Maddison & Tarlow , 2014 ; Bielik et al. , 2016 ; Parisotto et al. , 2017 ; Yin & Neubig , 2017 ; Rabinovich et al. , 2017 ) , which are syntactically correct by construction . In this work , we follow the AST generation approach . The key idea is to construct the AST a sequentially , by expanding one node at a time using production rules from the underlying programming language grammar . This simplifies the code generation task to a sequence of classification problems , in which an appropriate production rule has to be chosen based on the context information and the partial AST generated so far . In this work , we simplify the problem further — similar to Maddison & Tarlow ( 2014 ) ; Bielik et al . ( 2016 ) — by fixing the order of the sequence to always expand the left-most , bottom-most nonterminal node . Alg . 1 illustrates the common structure of AST-generating models . Then , the probability of generating a given AST a given some context c is p ( a | c ) = ∏ t p ( at | c , a < t ) , ( 1 ) where at is the production choice at step t and a < t the partial syntax tree generated before step t. Code Generation as Hole Completion We introduce the ExprGen task of filling in code within a hole of an otherwise existing program . This is similar , but not identical to the auto-completion function in a code editor , as we assume information about the following code as well and aim to generate whole expressions rather than single tokens . The ExprGen task also resembles program sketching ( Solar-Lezama , 2008 ) but we give no other ( formal ) specification other than the surrounding code . Concretely , we restrict ourselves to expressions that have Boolean , arithmetic or string type , or arrays of such types , excluding expressions of other types or expressions that use project-specific APIs . An example is shown in Fig . 1 . We picked this subset because it already has rich semantics that can require reasoning about the interplay of different variables , while it still only relies on few operators and does not require to solve the problem of open vocabularies of full programs , where an unbounded number of methods would need to be considered . In our setting , the context c is the pre-existing code around a hole for which we want to generate an expression . This also includes the set of variables v1 , . . . , v ` that are in scope at this point , which can be used to guide the decoding procedure ( Maddison & Tarlow , 2014 ) . Note , however , that our method is not restricted to code generation and can be easily extended to all other tasks and domains that can be captured by variations of Alg . 1 ( e.g . in NLP ) . 3 GRAPH DECODING FOR SOURCE CODE . To tackle the code generation task presented in the previous section , we have to make two design choices : ( a ) we need to find a way to encode the code context c , v1 , . . . , v ` and ( b ) we need to construct a model that can learn p ( at | c , a < t ) well . We do not investigate the question of encoding the context in this paper , and use two existing methods in our experiments in Sect . 5 . Both these encoders yield a distributed vector representation for the overall context , representations ht1 , . . . , htT for all tokens in the context , and separate representations for each of the in-scope variables v1 , . . . , v ` , summarizing how each variable is used in the context . This information can then be used in the generation process , which is the main contribution of our work and is described in this section . Overview Our decoder model follows the grammar-driven AST generation strategy of prior work as shown in Alg . 1 . The core difference is in how we compute the representation of the node to expand . Maddison & Tarlow ( 2014 ) construct it entirely from the representation of its parent in the AST using a log-bilinear model . Rabinovich et al . ( 2017 ) construct the representation of a node using the parents of the AST node but also found it helpful to take the relationship to the parent node ( e.g . “ condition of a while ” ) into account . Yin & Neubig ( 2017 ) on the other hand propose to take the last expansion step into account , which may have finished a subtree “ to the left ” . In practice , these additional relationships are usually encoded by using gated recurrent units with varying input sizes . We propose to generalize and unify these ideas using a graph to structure the flow of information in the model . Concretely , we use a variation of attribute grammars ( Knuth , 1967 ) from compiler theory to derive the structure of this graph . We associate each node in the AST with two fresh nodes representing inherited resp . synthesized information ( or attributes ) . Inherited information is derived from the context and parts of the AST that are already generated , whereas synthesized information can be viewed as a “ summary ” of a subtree . In classical compiler theory , inherited attributes usually contain information such as declared variables and their types ( to allow the compiler to check that only declared variables are used ) , whereas synthesized attributes carry information about a subtree “ to the right ” ( e.g. , which variables have been declared ) . Traditionally , to implement this , the language grammar has to be extended with explicit rules for deriving and synthesizing attributes . To transfer this idea to the deep learning domain , we represent attributes by distributed vector representations and train neural networks to learn how to compute attributes . Our method for getRepresentation from Alg . 1 thus factors into two parts : a deterministic procedure that turns a partial AST a < t into a graph by adding additional edges that encode attribute relationships , and a graph neural network that learns from this graph . Notation Formally , we represent programs as graphs where nodes u , v , . . . are either the AST nodes or their associated attribute nodes , and typed directed edges 〈u , τ , v〉 ∈ E connect the nodes according to the flow of information in the model . The edge types τ represent different syntactic or semantic relations in the information flow , discussed in detail below . We write Ev for the set of incoming edges into v. We also use functions like parent ( a , v ) and lastSibling ( a , v ) that look up and return nodes from the AST a ( e.g . resp . the parent node of v or the preceding AST sibling of v ) . Algorithm 2 Pseudocode for ComputeEdge Input : Partial AST a , node v 1 : Edge set E ← ∅ 2 : if v is inherited then 3 : E ← E ∪ { 〈parent ( a , v ) , Child , v〉 } 4 : if v is terminal node then 5 : E ← E ∪ { 〈lastToken ( a , v ) , NextToken , v〉 } 6 : if v is variable then 7 : E ← E ∪ { 〈lastUse ( a , v ) , NextUse , v〉 } 8 : if v is not first child then 9 : E ← E ∪ { 〈lastSibling ( a , v ) , NextSib , v〉 } 10 : else 11 : E ← E ∪ { 〈u , Parent , v〉 | u ∈ children ( a , v ) } 12 : E ← E ∪ { 〈inheritedAttr ( v ) , InhToSyn , v〉 } 13 : return E Example Consider the AST of the expression i - j shown in Fig . 2 ( annotated with attribute relationships ) constructed step by step by our model . The AST derivation using the programming language grammar is indicated by shaded backgrounds , nonterminal nodes are shown as rounded rectangles , and terminal nodes are shown as rectangles . We additionally show the variables given within the context as dashed rectangles at the bottom . First , the root node , Expr , was expanded using the production rule ( 1 ) : Expr =⇒ Expr - Expr . Then , its two nonterminal children were in turn expanded to the set of known variables using the produc- tion rule ( 2 ) : Expr =⇒ V , choosing i for the first variable and j for the second variable ( cf . below for details on picking variables ) . Attribute nodes are shown overlaying their corresponding AST nodes . For example , the root node is associated with its inherited attributes node 0 and with node 10 for its synthesized attributes . For simplicity , we use the same representation for inherited and synthesized attributes of terminal nodes . Edges in a < t We discuss the edges used in our neural attribute grammars ( NAG ) on our example below , and show them in Fig . 2 using different edge drawing styles for different edge types . Once a node is generated , the edges connecting this node can be deterministically added to a < t ( precisely defined in Alg . 2 ) . The list of different edge types used in our model is as follows : • Child ( red ) edges connect an inherited attribute node to the inherited attributes nodes of its children , as seen in the edges from node 0 . These are the connections in standard syntaxdriven decoders ( Maddison & Tarlow , 2014 ; Parisotto et al. , 2017 ; Yin & Neubig , 2017 ; Rabinovich et al. , 2017 ) . • Parent ( green ) edges connect a synthesized attribute node to the synthesized attribute node of its AST parent , as seen in the edges leading to node 10 . These are the additional connections used by the R3NN decoder introduced by Parisotto et al . ( 2017 ) . • NextSib ( black ) edges connect the synthesized attribute node to the inherited attribute node of its next sibling ( e.g . from node 5 to node 6 ) . These allow information about the synthesized attribute nodes from a fully generated subtree to flow to the next subtree . • NextUse ( orange ) edges connect the attribute nodes of a variable ( since variables are always terminal nodes , we do not distinguish inherited from synthesized attributes ) to their next use . Unlike Allamanis et al . ( 2018b ) , we do not perform a dataflow analysis , but instead just follow the lexical order . This can create edges from nodes of variables in the context c ( for example , from node 1 to 4 in Fig . 2 ) , or can connect AST leaf nodes that represent multiple uses of the same variable within the generated expressions . • NextToken ( blue ) edges connect a terminal node ( a token ) to the next token in the program text , for example between nodes 4 and 6 . • InhToSyn edges ( not shown in Fig . 2 ) connect the inherited attributes nodes to its synthe- sized attribute nodes . This is not strictly adding any information , but we found it to help with training . The panels of Fig . 2 show the timesteps at which the representations of particular attribute nodes are computed and added to the graph . For example , in the second step , the attributes for the terminal token i ( node 4 ) in Fig . 2 are computed from the inherited attributes of its AST parent Expr ( node 3 ) , the attributes of the last use of the variable i ( node 1 ) , and the node label i . In the third step , this computed attribute is used to compute the synthesized attributes of its AST parent Expr ( node 5 ) . Attribute Node Representations To compute the neural attribute representation hv of an attribute node v whose corresponding AST node is labeled with ` v , we first obtain its incoming edges using Alg . 2 and then use the state update function from Gated Graph Neural Networks ( GGNN ) ( Li et al. , 2016 ) . Thus , we take the attribute representations hui at edge sources ui , transform them according to the corresponding edge type ti using a learned function fti , aggregate them ( by elementwise summation ) and combine them with the learned embedding emb ( ` v ) of the node label ` v using a function g : hv = g ( emb ( ` v ) , ∑ 〈ui , ti , v〉∈Ev fti ( hui ) ) ( 2 ) In practice , we use a single linear layer for fti and implement g as a gated recurrent unit ( Cho et al. , 2014 ) . We compute node representations in such an order that all hui appearing on the right of ( 2 ) are already computed . This is possible as the graphs obtained by repeated application of Alg . 2 are directed acyclic graphs rooted in the inherited attribute node of the root node of the AST . We initialize the representation of the root inherited attribute to the representation returned by the encoder for the context information . Choosing Productions , Variables & Literals We can treat picking production rules as a simple classification problem over all valid production rules , masking out those choices that do not correspond to the currently considered nonterminal . For a nonterminal node v with label ` v and inherited attributes hv , we thus define pickProduction ( ` v , hv ) = argmaxP ( rule | ` v , hv ) = argmax [ e ( hv ) +m ` v ] . ( 3 ) Here , m ` v is a mask vector whose value is 0 for valid productions ` v ⇒ . . . and −∞ for all other productions . In practice , we implement e using a linear layer . Similarly , we pick variables from the set of variables V in scope using their representations hvvar ( initially the representation obtained from the context , and later the attribute representation of the last node in the graph in which they have been used ) by using a pointer network ( Vinyals et al. , 2015 ) . Concretely , to pick a variable at node v , we use learnable linear function k and define pickVariable ( V , hv ) = argmax var∈V P ( var | hv ) = argmax var∈V k ( hv , hvvar ) . ( 4 ) Note that since the model always picks a variable from the set of in-scope variables V , this generation model can never predict an unknown or out-of-scope variable . Finally , to generate literals , we combine a small vocabulary L of common literals observed in the training data and special UNK tokens for each type of literal with another pointer network that can copy one of the tokens t1 . . . tT from the context . Thus , to pick a literal at node v , we define pickLiteral ( V , hv ) = argmax lit∈L∪ { t1 ... tT } P ( lit | hv ) . ( 5 ) Note that this is the only operation that may produce an unknown token ( i.e . an UNK literal ) . In practice , we implement this by learning two functions sL and sc , such that sL ( hv ) produces a score for each token from the vocabulary and sc ( hv , hti ) computes a score for copying token ti from the context . By computing a softmax over all resulting values and normalizing it by summing up entries corresponding to the same constant , we can learn to approximate the desired P ( lit | hv ) . Training & Training Objective The different shapes and sizes of generated expressions complicate an efficient training regime . However , note that given a ground truth target tree , we can easily augment it with all additional edges according to Alg . 2 . Given that full graph , we can compute a propagation schedule ( intuitively , a topological ordering of the nodes in the graph , starting in the root node ) that allows to repeatedly apply ( 2 ) to obtain representations for all nodes in the graph . By representing a batch of graphs as one large ( sparse ) graph with many disconnected components , similar to Allamanis et al . ( 2018b ) , we can train our graph neural network efficiently . We have released the code for this on https : //github.com/Microsoft/graph-based-code-modelling . Our training procedure thus combines an encoder ( cf . Sect . 5 ) , whose output is used to initialize the representation of the root and context variable nodes in our augmented syntax graph , the sequential graph propagation procedure described above , and the decoder choice functions ( 3 ) and ( 4 ) . We train the system end-to-end using a maximum likelihood objective without pre-trained components . Additional Improvements We extend ( 3 ) with an attention mechanism ( Bahdanau et al. , 2014 ; Luong et al. , 2015 ) that uses the state hv of the currently expanded node v as a key and the context token representations ht1 , . . . , htT as memories . Experimentally , we found that extending Eqs . 4 , 5 similarly did not improve results , probably due to the fact that they already are highly dependent on the context information . Following Rabinovich et al . ( 2017 ) , we provide additional information for Child edges . To allow this , we change our setup so that some edge types also require an additional label , which is used when computing the messages sent between different nodes in the graph . Concretely , we extend ( 2 ) by considering sets of unlabeled edges Ev and labeled edges E ` v : hv = g ( emb ( ` v ) , ∑ ( ui , ti , v ) ∈Ev fti ( hui ) + ∑ ( ui , ti , ` i , v ) ∈E ` v fti ( hui , embe ( ` i ) ) ) ( 6 ) Thus for labeled edge types , fti takes two inputs and we additionally introduce a learnable embedding for the edge labels . In our experiments , we found it useful to label Child with tuples consisting of the chosen production and the index of the child , i.e. , in Fig . 2 , we would label the edge from 0 to 3 with ( 2 , 0 ) , the edge from 0 to 6 with ( 2 , 1 ) , etc . Furthermore , we have extended pickProduction to also take the information about available variables into account . Intuitively , this is useful in cases of productions such as Expr =⇒ Expr.Length , which can only be used in a well-typed derivation if an array-typed variable is available . Thus , we extend e ( hv ) from ( 3 ) to additionally take the representation of all variables in scope into account , i.e. , e ( hv , r ( { hvvar | var ∈ V } ) ) , where we have implemented r as a max pooling operation .

Summary: The paper introduces a 'code generation as hole completion' task and associated dataset, ExprGen. The authors proposed a novel extension of AST code generation which uses what they call Neural Attribute Grammars. They show the proposed method does well on this task, compared to ablations of their model (which are similar to previous AST approaches).