Prediction of Types in Python with Pre-trained Graph Neural Networks

The application of Graph Neural Networks for pre-training models for source code is not well studied. We experimented with pre-training a Graph Neural Network model for Python on tasks of Name Prediction and Edge Prediction. Then, we used pre-trained weights to initialize a model for variable type prediction. Our preliminary results suggest that pre-training on these tasks brings neither improvements in type prediction performance nor training dynamics. Possible ways to fix this are discussed in the concluding section of the paper. Additionally, we performed an ablation study to see whether type prediction is overreliant on some parts of the graph. Results suggest, that type prediction model does not significantly rely on obvious shortcuts and could be a useful proxy for evaluating pre-trained graph embeddings.