Axon: A Language for Dynamic Shapes in Deep Learning Graphs

TL;DR

Axon introduces a functional language with shape inference capabilities for deep learning graphs, enabling automatic tensor shape deduction, reducing programmer effort, and improving error detection in complex models.

Contribution

It presents a novel shape inference algorithm based on constraint solving, integrating shape inference into a functional language for deep learning.

Findings

Enables automatic inference of tensor shapes in complex graphs

Reduces need for explicit shape annotations by programmers

Improves error detection through shape mismatch identification

Abstract

Axon is a language that enables shape and rank inference for tensors in a Deep Learning graphs. It aims to make shapes implicit and inferred, in a similar manner to how types are implicit and inferred in many functional programming languages. Tensor dimensions are represented by expressions consisting of symbolic variables, constants, and arithmetic operators. Tensor shapes can be expressed as either a sequence of these dimension expressions, as a symbolic variable, or as an appending of other shapes. This allows complex constraints on shapes to be expressed. Axon is functional in style, with a type system similar in to Standard ML, extended to include shape information. It provides a suite of built in operators over tensors, including pointwise arithmetic operators, maps, reduction, loops and user defined functions. We describe a shape inference algorithm based on constraint solving which infers information about shapes, from both shape information provided by the programmer and the structure of the program. This allows fully automatic inference of the shapes of tensors for complex Deep Learning graphs. This approach reduces programmer effort when specifying graphs, as tensor shapes are not explicit, allows composition of Deep Learning graphs while maintaining input and output tensor shape compatibility, and aids in automated error detection by identifying shape mismatches at runtime.