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: This paper proposed to use Graph Neural Networks (GNN) to do type inference for dynamically typed languages. The key technique is to construct a type dependency graph and infer the type on top of it. The type dependency graph contains edges specifying hard constraints derived from the static analysis, as well as soft relationships specified by humans. Experiments on type predictions for TypeScript have shown better performance than the previous methods, with or without user specified types.