Source-linked AI summary
LambdaNet: Probabilistic Type Inference using Graph Neural Networks
Jiayi Wei, Maruth Goyal, Greg Durrett, Isil Dillig
TL;DR
Type annotations can support maintainability, compile-time error catching, and code completion, but inference and manual annotation are difficult, and prior TypeScript systems have limited type coverage and accuracy. LAMBDANET uses a graph neural network on a statically extracted type dependency graph to predict library and user-defined types. It reports higher library-type accuracy than DeepTyper and can predict user-defined types not seen during training.
Problem
Prior TypeScript inference systems are limited to training-seen types and have relatively low accuracy, while annotations are difficult to infer automatically or add by hand.
Source: §1 · p. 1Method
Results
Takeaways & Limitations
LAMBDANET combines explicit program analysis with graph neural networks to predict library and user-defined types, including types absent from training.
Source: §7 · p. 9The prediction scope excludes polymorphic and function types, and structured-type prediction is left for future work.
Source: §2.1 · p. 3
Abstract
from arXiv · showhide
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 cannot 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
Type annotations can improve maintainability, compile-time error catching, and code completion, but inference and manual annotation remain difficult. LAMBDANET uses a graph neural network over a type dependency graph and outperforms prior work on library types while predicting user-defined types.
1 INTRODUCTION
- Dynamic languages lack static type systems, while compilers cannot infer types using standard algorithms and manual annotation is tedious and error-prone.The paper frames gradual typing as a way to support static checking when annotations are available.Source: §1 · p. 1
- Prior TypeScript inference systems are limited to types seen during training, achieve 56.9% accuracy for primitive/library types, and can assign inconsistent types to occurrences of one variable.The cited state of the art is reported at 56.9% accuracy for primitive/library types.Source: §1 · p. 1
- LAMBDANET represents programs as type dependency graphs that encode logical constraints alongside contextual hints from names and variable usage.A graph neural network processes the graph to predict types, including user-defined types absent from training.Source: §1 · p. 2
- 75.6% top1 accuracy on library types, compared with DeepTyper’s 61.5%, and around 64.2% overall accuracy including user-defined types.The reported overall accuracy is 55.2% (absolute) higher than the TypeScript compiler.Source: §1 · p. 2
- The paper’s contributions include GNN embeddings for type variables, pointer-network-like prediction of user-defined types, and evaluation on hundreds of real-world TypeScript projects.The paper reports significant improvement over prior work.Source: §1 · p. 2
2 MOTIVATING EXAMPLE AND PROBLEM SETTING
The motivating example illustrates why type inference needs both constraints and contextual clues: operators restrict possible types, while names and usage can help distinguish among structurally plausible types. The task predicts a probability distribution over library and project-defined types, with structured-type prediction outside the stated scope.
2 MOTIVATING EXAMPLE AND PROBLEM SETTING
- Operator use imposes hard constraints: for example, concat requires x and y to support the operation, ruling out boolean while allowing types such as string, array, or Tensor.These constraints motivate including typing constraints in the model.Source: §2 · p. 2
- A variable name like network and the class name MyNetwork provide a contextual clue that can help distinguish MyNetwork from other classes with a time field.The example also propagates the number return type of readNumber to MyNetwork’s time field, assuming that library return type is known.Source: §2 · pp. 2–3
- The type dependency graph represents predicates relevant to inference, extracted by static analysis rather than relying only on token sequences, syntax trees, or control-flow graphs.Figure 2 represents the unannotated program with type variables, including variables introduced for intermediate expressions.Source: §2 · p. 3, Figure 2 · p. 3
2.1 PROBLEM SETTING
- The model predicts a probability distribution for each missing annotation over fixed library types and classes or interfaces declared in the project.Figure 3 illustrates subtype, object-member, identifier, and usage relationships in the example graph.Source: §2.1 · p. 3, Figure 3 · p. 3
- The prediction scope excludes polymorphic and function types, collapsing generic instantiations and function signatures into broader types.Predicting structured types is left for future work.Source: §2.1 · p. 3
3 TYPE DEPENDENCY GRAPH
The type dependency graph is a hypergraph built through static analysis of an intermediate representation, with type variables as nodes and labeled hyperedges encoding predicates. Its edges capture both logical constraints and contextual information, including variable names, name similarity, and object-usage relationships.
3 TYPE DEPENDENCY GRAPH
- Static analysis of an intermediate representation assigns a unique variable to each program sub-expression and extracts a hypergraph whose nodes are type variables and whose labeled hyperedges encode their relationships.Each hyperedge corresponds to a predicate, partitioned into Logical constraints and Contextual hints.Source: §3 · p. 4
- Logical hyperedges encode hard constraints, while Contextual hyperedges provide hints from variable, function, and class names.The example includes subtype and object-member relationships as logical predicates, and an identifier-associated variable as a contextual clue.Source: §3 · p. 4
- NameSimilar edges connect variables whose tokenized names overlap, capturing correspondences such as a MyNetwork instance named network.Identifier names are split using underscores and camel case.Source: §3 · p. 4
- Usage edges connect an accessed object and expression to classes with a matching member name, allowing the GNN to pass information along those relationships.The model uses a special attention mechanism on usage edges.Source: §3 · p. 4
4 NEURAL ARCHITECTURE
LambdaNet’s neural architecture uses typed-graph message passing to create embeddings for type variables, then uses a pointer-network-like layer to assign probabilities to candidate types. Its propagation and initialization are tailored to different hyperedge types and to known versus unknown types.
4 NEURAL ARCHITECTURE
- A graph neural network produces type-variable embeddings, which a pointer-network-like layer compares with candidate type embeddings to distribute probability over type assignments.Library types use trainable vectors, while user-defined types use their graph-node embeddings.Source: §4 · pp. 5–6
- Across K message-passing rounds, typed hyperedge operators send messages and aggregation combines incoming messages to update each node embedding.The neighborhood function and edge type determine which message operation is used.Source: §4 · pp. 5–6
- Known-type constant nodes keep a trainable type-specific embedding, while variable nodes begin with the same generic trainable vector.Constant embeddings remain unchanged during GNN iterations; variable embeddings are updated.Source: §4 · p. 5
- FIXED edges use concatenated arguments and an MLP, NARY edges send position-labeled messages, and NPAIRS edges use dot-product attention.For NARY edges, messages to one argument do not depend on the other arguments’ embeddings.Source: §4 · p. 5
- Aggregation uses a graph-attention variant with dot-product attention weights and a LeakyReLU slope of 0.2.Weights are normalized with softmax; the attention calculation uses trainable matrices.Source: §4 · p. 6
- Identifier embeddings split names by camel case and underscores, while rare tokens are randomly mapped among 51 unknown-token vectors.The random mapping is reconstructed each GNN run, and embeddings are trained end-to-end.Source: §4 · p. 6
5 EVALUATION
Across 300 real-world TypeScript projects, LAMBDANET is evaluated against prior systems, on user-defined-type prediction, and through component ablations. It outperforms the reported baselines, while the ablations show that message passing and both logical and contextual edges contribute to performance.
5 EVALUATION
- Evaluation uses 300 GitHub projects, with 60 held out for testing, and tests annotations manually added by developers.Training labels are augmented using TypeScript compiler inference, excluding labels assigned any; predictions cover project-defined classes/interfaces and the 100 most common library types.Source: §5 · pp. 6–7
5.2 PREDICTING USER-DEFINED TYPES
- 64.2% top-1 and 84.5% top-5 overall accuracy significantly outperform both the TypeScript compiler and SIMILARNAME baselines.Table 2 reports accuracy for user-defined and library types individually as well as overall.Source: Table 2 · p. 8, §5.2 · p. 8
5.3 ABLATION STUDY
- Table 3 shows accuracy improving through six message-passing iterations; removing contextual or logical edges drops overall accuracy by 14.5% or 25.8%, respectively.The iteration trend is reported as evidence that the network learns inference over long distances.Source: Table 3 · p. 8, §5.3 · pp. 8–9
5.1 COMPARISON WITH DEEPTYPER
- LAMBDANET beats DeepTyper by 14.1% absolute declaration-level accuracy and 9.6% occurrence-level accuracy.The comparison fixes both systems’ prediction spaces to library types.Source: §5.1 · pp. 7–8
5.4 COMPARISON WITH JSNICE
- On 41 sampled top-level functions, LAMBDANET correctly predicts 77 of 107 annotations, compared with 48 for JSNice.The benchmark excludes functions with non-library-type parameters and includes dependency definitions.Source: §5.4 · p. 9
6 RELATED WORK
Prior type-inference work includes structured models, deep-learning models over sequences or syntax trees, and approaches using documentation strings. The related work also contrasts program-graph applications of GNNs and methods for predicting labels outside a fixed vocabulary.
Type Inference using Statistical Methods
- Earlier type predictors use structured inference, sequence or syntax-tree deep learning, or documentation strings, but are limited to a fixed type vocabulary.The cited systems target JavaScript, Python, and TypeScript.Source: §6 · p. 9
Graph Embedding of Programs
- Prior GNN applications include predicting variable names and misuses from C# programs, encoding formulas for theorem proving, and using attention to compute node embeddings.LAMBDANET instead uses attention to compute certain messages from node embeddings.Source: §6 · p. 9
Predicting from an Open Vocabulary
- Open-vocabulary methods often compare input and label embeddings; LAMBDANET derives type encodings during the input’s forward pass, as in unknown-word handling for machine translation.The passage identifies this input-derived encoding as what makes its approach a pointer network.Source: §6 · p. 9
7 CONCLUSIONS
LAMBDANET combines explicit program analysis with graph neural networks to predict library and previously unseen user-defined types. Its architecture simplifies function and generic types, while structured-type prediction and hard consistency constraints remain future work.
7 CONCLUSIONS
- LAMBDANET combines explicit program analysis with graph neural networks and can predict user-defined types unseen during training.Ablation studies demonstrate the usefulness of logical and contextual hyperedges.Source: §7 · p. 9
- The architecture simplifies function and generic types by collapsing them into their non-generic counterparts.Source: §7 · pp. 9–10
- Future work includes predicting structured types and enforcing hard constraints so resulting type assignments are guaranteed consistent.Source: §7 · p. 10