Dynamic Inference of Likely Symbolic Tensor Shapes in Python Machine Learning Programs
Daniel Zheng, Koushik Sen · 2024
In machine learning programs, it is often tedious to annotate the dimensions of shapes of various tensors that get created during execution. We present a dynamic likely tensor shape inference analysis, called ShapeIt, that annotates the dimensions of shapes of tensor expressions with symbolic dimension values and establishes the symbolic relationships among those dimensions. Such annotations can be used to understand the machine learning code written in popular frameworks, such as PyTorch and JAX, and to find bugs related to tensor shape mismatch. We have implemented ShapeIt on top of a novel dynamic analysis framework for Python, called Pynsy, which works by instrumenting Python bytecode on the fly. Our evaluation of ShapeIt on several tensor programs illustrates that ShapeIt could effectively infer symbolic shapes and their relationships for various neural network programs with low runtime overhead.