Gradual Tensor Shape Checking

TL;DR

This paper introduces a gradual, type-based approach for tensor shape checking in deep learning programs, combining static inference with dynamic checks to reduce bugs related to shape mismatches.

Contribution

It presents a novel gradual typing system for tensor shape inference that balances static analysis and dynamic verification, formalized and implemented in a prototype tool.

Findings

Effective detection of shape mismatches in neural networks

Successful implementation of a shape checking tool

Improved reliability of deep learning programs

Abstract

Tensor shape mismatch is a common source of bugs in deep learning programs. We propose a new type-based approach to detect tensor shape mismatches. One of the main features of our approach is the best-effort shape inference. As the tensor shape inference problem is undecidable in general, we allow static type/shape inference to be performed only in a best-effort manner. If the static inference cannot guarantee the absence of the shape inconsistencies, dynamic checks are inserted into the program. Another main feature is gradual typing, where users can improve the precision of the inference by adding appropriate type annotations to the program. We formalize our approach and prove that it satisfies the criteria of gradual typing proposed by Siek et al. in 2015. We have implemented a prototype shape checking tool based on our approach and evaluated its effectiveness by applying it to some deep neural network programs.