LambdaNet: Probabilistic Type Inference using Graph Neural Networks

As gradual typing becomes increasingly popular in languages like Python and TypeScript , there is a growing need to infer type annotations automatically . While type annotations help with tasks like code completion and static error catching , these annotations can not be fully determined by compilers and are tedious to annotate by hand . This paper proposes a probabilistic type inference scheme for TypeScript based on a graph neural network . Our approach first uses lightweight source code analysis to generate a program abstraction called a type dependency graph , which links type variables with logical constraints as well as name and usage information . Given this program abstraction , we then use a graph neural network to propagate information between related type variables and eventually make type predictions . Our neural architecture can predict both standard types , like number or string , as well as user-defined types that have not been encountered during training . Our experimental results show that our approach outperforms prior work in this space by 14 % ( absolute ) on library types , while having the ability to make type predictions that are out of scope for existing techniques . 1 INTRODUCTION . Dynamically typed languages like Python , Ruby , and Javascript have gained enormous popularity over the last decade , yet their lack of a static type system comes with certain disadvantages in terms of maintainability ( Hanenberg et al. , 2013 ) , the ability to catch errors at compile time , and code completion support ( Gao et al. , 2017 ) . Gradual typing can address these shortcomings : program variables have optional type annotations so that the type system can perform static type checking whenever possible ( Siek & Taha , 2007 ; Chung et al. , 2018 ) . Support for gradual typing now exists in many popular programming languages ( Bierman et al. , 2014 ; Vitousek et al. , 2014 ) , but due to their heavy use of dynamic language constructs and the absence of principal types ( Ancona & Zucca , 2004 ) , compilers can not perform type inference using standard algorithms from the programming languages community ( Bierman et al. , 2014 ; Traytel et al. , 2011 ; Pierce & Turner , 2000 ) , and manually adding type annotations to existing codebases is a tedious and error-prone task . As a result , legacy programs in these languages do not reap all the benefits of gradual typing . To reduce the human effort involved in transitioning from untyped to statically typed code , this work focuses on a learning-based approach to automatically inferring likely type annotations for untyped ( or partially typed ) codebases . Specifically , we target TypeScript , a gradually-typed variant of Javascript for which plenty of training data is available in terms of type-annotated programs . While there has been some prior work on inferring type annotations for TypeScript using machine learning ( Hellendoorn et al. , 2018 ; Raychev et al. , 2015 ) , prior work in this space has several shortcomings . First , inference is restricted to a finite dictionary of types that have been observed during training time—i.e. , they can not predict any user-defined data types . Second , even without considering user-defined types , the accuracy of these systems is relatively low , with the current state-of-theart achieving 56.9 % accuracy for primitive/library types ( Hellendoorn et al. , 2018 ) . Finally , these techniques can produce inconsistent results in that they may predict different types for different token-level occurrences of the same variable . In this paper , we propose a new probabilistic type inference algorithm for TypeScript to address these shortcomings using a graph neural network architecture ( GNN ) ( Veličković et al. , 2018 ; Li et al. , 2016 ; Mou et al. , 2016 ) . Our method uses lightweight source code analysis to transform the program into a new representation called a type dependency graph , where nodes represent type variables and labeled hyperedges encode relationships between them . In addition to expressing logical constraints ( e.g. , subtyping relations ) as in traditional type inference , a type dependency graph also incorporates contextual hints involving naming and variable usage . Given such a type dependency graph , our approach uses a GNN to compute a vector embedding for each type variable and then performs type prediction using a pointer-network-like architecture ( Vinyals et al. , 2015 ) . The graph neural network itself requires handling a variety of hyperedge types—some with variable numbers of arguments—for which we define appropriate graph propagation operators . Our prediction layer compares the vector embedding of a type variable with vector representations of candidate types , allowing us to flexibly handle user-defined types that have not been observed during training . Moreover , our model predicts consistent type assignments by construction because it makes variable-level rather than token-level predictions . We implemented our new architecture as a tool called LAMBDANET and evaluated its performance on real-world TypeScript projects from Github . When only predicting library types , LAMBDANET has a top1 accuracy of 75.6 % , achieving a significant improvement over DeepTyper ( 61.5 % ) . In terms of overall accuracy ( including user-defined types ) , LAMBDANET achieves a top1 accuracy of around 64.2 % , which is 55.2 % ( absolute ) higher than the TypeScript compiler . Contributions . This paper makes the following contributions : ( 1 ) We propose a probabilistic type inference algorithm for TypeScript that uses deep learning to make predictions from the type dependency graph representation of the program . ( 2 ) We describe a technique for computing vector embeddings of type variables using GNNs and propose a pointer-network-like method to predict user-defined types . ( 3 ) We experimentally evaluate our approach on hundreds of real-world TypeScript projects and show that our method significantly improves upon prior work . 2 MOTIVATING EXAMPLE AND PROBLEM SETTING . Figure 1 shows a ( type-annotated ) TypeScript program . Our goal in this work is to infer the types shown in the figure , given an unannotated version of this code . We now justify various aspects of our solution using this example . Typing constraints . The use of certain functions/operators in Figure 1 imposes hard constraints on the types that can be assigned to program variables . For example , in the forward function , variables x , y must be assigned a type that supports a concat operation ; hence , x , y could have types like string , array , or Tensor , but not , for example , boolean . This observation motivates us to incorporate typing constraints into our model . Contextual hints . Typing constraints are not always sufficient for determining the intended type of a variable . For example , for variable network in function restore , the typing constraints require network ’ s type to be a class with a field called time , but there can be many classes that have such an attribute ( e.g. , Date ) . However , the similarity between the variable name network and the class name MyNetwork hints that network might have type MyNetwork . Based on this belief , we can further propagate the return type of the library function readNumber ( assuming we know it is number ) to infer that the type of the time field in MyNetwork is likely to be number . Need for type dependency graph . There are many ways to view programs—e.g. , as token sequences , abstract syntax trees , control flow graphs , etc . However , none of these representations is particularly helpful for inferring the most likely type annotations . Thus , our method uses static analysis to infer a set of predicates that are relevant to the type inference problem and represents these predicates using a program abstraction called the type dependency graph . Handling user-defined types . As mentioned in Section 1 , prior techniques can only predict types seen during training . However , the code from Figure 1 defines its own class called MyNetwork and later uses a variables of type MyNetwork in the restore method . A successful model for this task therefore must dynamically make inferences about user-defined types based on their definitions . 2.1 PROBLEM SETTING . Our goal is to train a type inference model that can take as input an entirely ( or partially ) unannotated TypeScript project g and output a probability distribution of types for each missing annotation . The prediction space is Y ( g ) = Ylib ∪ Yuser ( g ) , where Yuser ( g ) is the set of all user-defined types ( classes/interfaces ) declared within g , and Ylib is a fixed set of commonly-used library types . Following prior work in this space ( Hellendoorn et al. , 2018 ; Raychev et al. , 2015 ; Xu et al. , 2016 ) , we limit the scope of our prediction to non-polymorphic and non-function types . That is , we do not distinguish between types such as List < T > , List < number > , List < string > etc. , and consider them all to be of type List . Similarly , we also collapse function types like number→ string and string→ string into a single type called Function . We leave the extension of predicting structured types as future work . 3 TYPE DEPENDENCY GRAPH . A type dependency graph G = ( N , E ) is a hypergraph where nodes N represent type variables and labeled hyperedges E encode relationships between them . We extract the type dependency graph of a given TypeScript program by performing static analysis on an intermediate representation of its source code , which allows us to associate a unique variable with each program sub-expression . As an illustration , Figure 2 shows the intermediate representation of the code from Figure 1 . Intuitively , a type dependency graph encodes properties of type variables as well as relationships between them . Each hyperedge corresponds to one of the predicates shown in Table 1 . We partition our predicates ( i.e. , hyperedges ) into two classes , namely Logical and Contextual , where the former category can be viewed as imposing hard constraints on type variables and the latter category encodes useful hints extracted from names of variables , functions , and classes . Figure 3 shows some of the hyperedges in the type dependency graph G extracted from the intermediate representation in Figure 2 . As shown in Figure 3 ( A ) , our analysis extracts a predicate Subtype ( τ13 , τ5 ) from this code because the type variable associated with the returned expression v4 must be a subtype of the enclosing function ’ s return type . Similarly , as shown in Figure 3 ( B ) , our analysis extracts a predicate Objectname , time , forward ( τ8 , τ1 , τ2 , τ9 ) because τ8 is an object type whose name , time , and forward members are associated with type variables τ1 , τ2 , τ9 , respectively . In contrast to the Subtype and Object predicates that impose hard constraints on type variables , the next two hyperedges shown in Figure 3 encode contextual clues obtained from variable names . Figure 3 ( C ) indicates that type variable τ14 is associated with an expression named restore . While this kind of naming information is invisible to TypeScript ’ s structural type system ( ? ) , it serves as a useful input feature for our GNN architecture described in Section 4 . In addition to storing the unique variable name associated with each type variable , the type dependency graph also encodes similarity between variable and class names . The names of many program variables mimic their types : for example , instances of a class called MyNetwork might often be called network or network1 . To capture this correspondence , our type dependency graph also contains a hyperedge called NameSimilar that connects type variables α and β if their corresponding tokenized names have a non-empty intersection.1 As shown in Table 1 , there is a final type of hyperedge called Usage that facilitates type inference of object types . In particular , if there is an object access var y = x.l , we extract the predicate Usagel ( ( τx , τy ) , ( α1 , β1 ) , . . . , ( αk , βk ) ) to connect x and y ’ s type variables with all classes αi that contain an attribute/method βi whose name is l. Figure 3 shows a Usage hyperedge extracted from the code in Figure 2 . As we will see in the next section , our GNN architecture utilizes a special attention mechanism to pass information along these usage edges . 1During tokenization , we split identifier names into tokens based on underscores and camel case naming . More complex schemes are possible , but we found this simple method to be effective .

Summary: A method to predict likely type of program variables in TypeScript is presented. It consists of a translation of a program's type constraints and defined objects into a (hyper)graph, and a specialised neural message passing architecture to learn from the generated graphs. Experiments show that the method substantially outperforms sound typing in the TypeScript compiler, as well as a recent method based on deep neural networks.