A Static Analyzer for Detecting Tensor Shape Errors in Deep Neural Network Training Code

TL;DR

PyTea is a static analysis tool that efficiently detects tensor shape errors in PyTorch code by analyzing all possible execution paths and tensor constraints, preventing costly runtime errors during neural network training.

Contribution

This paper introduces PyTea, a static analyzer capable of accurately detecting tensor shape errors in PyTorch code with high scalability and precision.

Findings

PyTea detects tensor shape errors within seconds in real-world PyTorch projects.

The analyzer effectively handles complex execution paths with minimal path explosion.

PyTea successfully identifies shape mismatches in code snippets from StackOverflow and official repositories.

Abstract

We present an automatic static analyzer PyTea that detects tensor-shape errors in PyTorch code. The tensor-shape error is critical in the deep neural net code; much of the training cost and intermediate results are to be lost once a tensor shape mismatch occurs in the midst of the training phase. Given the input PyTorch source, PyTea statically traces every possible execution path, collects tensor shape constraints required by the tensor operation sequence of the path, and decides if the constraints are unsatisfiable (hence a shape error can occur). PyTea's scalability and precision hinges on the characteristics of real-world PyTorch applications: the number of execution paths after PyTea's conservative pruning rarely explodes and loops are simple enough to be circumscribed by our symbolic abstraction. We tested PyTea against the projects in the official PyTorch repository and some tensor-error code questioned in the StackOverflow. PyTea successfully detects tensor shape errors in these codes, each within a few seconds.