Axon A Language for Dynamic Shapes in Deep Learning Graphs_2

2025-04-27 0 0 171.79KB 9 页 10玖币

侵权投诉

arXiv:2210.02374v1 [cs.PL] 5 Oct 2022

Axon: A Language for Dynamic Shapes

in Deep Learning Graphs

Alexander Collins

NVIDIA

acollins@nvidia.com

Vinod Grover

NVIDIA

vgrover@nvidia.com

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 lan-

guages. Tensor dimensions are represented by expressions

consisting of symbolic variables, constants, and arithmetic

operators. Tensor shapes can be expressed as either a se-

quence of these dimension expressions, as a symbolic vari-

able, or as an appending of other shapes. This allows com-

plex 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 deﬁned functions.

We describe a shape inference algorithm based on con-

straint solving which infers information about shapes, from

both shape information provided by the programmer and

the structure of the program. This allows fully automatic in-

ference of the shapes of tensors for complex Deep Learning

graphs.

This approach reduces programmer eﬀort when specify-

ing graphs, as tensor shapes are not explicit, allows compo-

sition of Deep Learning graphs while maintaining input and

output tensor shape compatibility, and aids in automated er-

ror detection by identifying shape mismatches at runtime.

1 Introduction

Deep Learning models can be viewed as constrained func-

tional programs on tensor domains, which only permit side

eﬀects or updates for certain types of models, and usually

only during training. Tensors have a type and a shape: they

are rectangular domains of simple element types. Requiring

the programmer to deal with these shapes can add signiﬁ-

cant complexity to a language.

This paper describes shape and rank inference for tensors

in a language we call Axon, which aims to make shapes im-

plicit and inferred, in a similar manner to how types are

implicit and inferred in many functional programming lan-

guages. Axon allows the individual dimensions of a tensor

to be expressions consisting of symbolic variables, constants,

and arithmetic operators, allowing complex constraints on

shapes to be expressed. Furthermore, shapes can be expressed

as either a sequence of these dimension expressions, as a

symbolic variable, or as an appending the dimensions of sev-

eral other shapes. This allows rank inference to be expressed

using this shape appending.

An inference algorithm is also presented which infers in-

formation about shapes, from both shape information pro-

vided by the programmer and the structure of the program,

via a set of rules for built-in operators. Our system allows

the shapes involved in complex Deep Learning problems to

be automatically inferred without need for the programmer

to express the shape constraints or give concrete shapes

which are potentially unknown until the shape of the in-

puts to the program are known. Shape information can be

inferred from known shapes, or partial shapes, of program

inputs and from the structure of the program itself. Auto-

mated error detection also aids programmers by identifying

shape mismatches between the allowed shapes to an opera-

tion and inferred shapes.

The rest of the paper is structured as follows. Section 2

outlines the Axon language. Section 3 presents the syntax of

shape expressions, section 4 describes how standard Hindley-

Milner type inference [4, 7] is used to generate a set of shape

constraints for the program and section 5 describes the algo-

rithm used to solve sets of these shape constraints. Section 6

provides examples of shape inference in action and section 7

provides larger examples for commonly seen Deep Learn-

ing graphs. Section 8 discusses related work and section 9

presents our conclusions.

This paper makes the following contributions:

•A taxonomy of the kinds of shapes encountered in

Deep Learning graphs, which we use to guide the de-

sign of the language.

•A functional programming language, which we call

Axon. It allows a programmer to specify symbolic shapes

for input and output tensors for a graph, which permit

arithmetic expressions for the dimensions, and allows

for rank inference by expressing shapes as the compo-

sition of sub-shapes.

•A constraint based solver for inferring shape informa-

tion throughout a program, based on a set of shape

inference rules. Initial constraints are generated based

on the structure of a graph, and the rules are used to

reduce these constraints to a ﬁxed point, where either

all shapes are known rank with constant dimensions,

or are partially unknown allowing for runtime vari-

able shape.

2 The Axon Language

The goal of Axon is to provide a language for easily writ-

ing Deep Learning graphs, in a target agnostic manner. It

provides a high-level functional language for representing

Deep Learning graphs, through the use of a suite of built-in

deep learning operators. The Axon compiler includes a type

inference algorithm, so that explicit types do not need to be

speciﬁed by the programmer. Furthermore, an inference al-

gorithm for shapes also removes the need for explicit tensor

shapes to be speciﬁed by the programmer.

A simple graph that compute the pointwise sum and prod-

uct of two inputs can be expressed in Axon as follows:

f = fn ( x , y) {

a = x + y;

b = x * y;

a , b

}

This program describes a graph called 𝑓which takes two

inputs 𝑥and 𝑦and produces two outputs 𝑎and 𝑏. In Axon,

all variables are tensors. These have an element type (with

the usual primitive types including ﬁxed size integers and

ﬂoating point data types, and a shape.

Axon allows program decomposition into sub-graphs as

in the following example:

g = fn ( a , b) {

a + b

};

f = fn ( a , b) {

g( a , b) + a;

}

This allows a program to be decomposed into smaller mod-

ular parts and reused. The type and shape inference algo-

rithm allow for polymorphic types and shapes. For example,

in the above graph, if 𝑔were called from multiple places it

could be called with diﬀerent types and shapes.

Axon includes a suite of operators for performing Deep

Learning computations. These include matrix multiplication,

convolution, point-wise arithmetic operators, loops and more.

The language includes a type speciﬁcation for each built-

in, which is used to infer types and shapes throughout the

graph. Examples of these are given in Section 3.3.

3 Tensor Shapes

In this section we describe the shapes of tensors in the Axon

language. Section 3.1 outlines a general taxonomy of shapes

we use to guide our design and discussion. These are a gen-

eral overview of the sorts of shapes that appear in Deep

Learning graphs. In Section 3.2 we give the concrete syn-

tax of shapes in the Axon. Examples of shape constraints

are given in Section 3.3, which show how shape constraints

are generated from the syntactic structure of a graph. The

algorithm for then solving these constraints is given later,

in Section 5

3.1 Taxonomy of Shapes

Here we specify the general taxonomy of shapes we use to

guide the design of our language. This taxonomy catego-

rizes the possible tensor shapes based on their properties

and constraints.

•Concrete shapes

These are shapes of statically known constant rank

whose dimensions are all statically known constants.

For example [] denotes a scalar, [256] denotes a one

dimensional tensor of 256 elements, and [1024, 256]

denotes a two dimension tensor.

•Symbolic dimensions

These extend concrete shapes by allowing individual

dimensions to be symbolic expressions. The rank of

the shape is still statically known.

For example [x,y,16] is a three dimensional tensor whose

inner two dimensions are symbolic variables. At run-

time, xand ywill be known, but at compile time their

value is unknown.

[x*2,y,16] is the shape that results from the concate-

nation of two tensors with the previous shape. Here

the outermost dimension is a symbolic expression.

•Symbolic

These extend symbolic dimension shapes by allowing

the for unknown rank.

For example s @ [x,y] denotes a tensor with at least

two dimensions. sis an arbitrary shape, and the @op-

erator is used to append sto the shape [x,y].

•Runtime dimensions

Here the rank is static, but the individual dimensions

can depend on the computation.

•Runtime

These extend runtime dimensions shapes by allowing

the rank to also depend on computation.

•Variable

Dimension values diﬀer across the tensor. This allows,

for example, ragged batches to be represented.

3.2 Syntax of Shapes

Figure 1 gives the syntax of shapes in Axon. The syntax of

shapes is too complex for a succinct grammar, therefore the

grammar given in Figure 1 generates some invalid shapes,

which are further constrained by the shape construction rules

given in Figure 2.

Shapes can be ﬁxed, or variable, according to the dimen-

sion expressions they contain. The syntax of these expres-

sions includes positive integer constants, symbolic variables

and arithmetic operations and a “variable dimension” de-

noted *. The *dimension expression (with no arguments)

denotes a variable dimension, i.e. the number of elements

in this dimension varies across the other dimensions in the

tensor. This allows, for example, batches of variable length

文档加载中……请稍候！
如果长时间未打开，您也可以点击刷新试试。

下载文档到电脑，查找使用更方便

10 玖币 0人已下载

立即下载

摘要：

arXiv:2210.02374v1[cs.PL]5Oct2022Axon:ALanguageforDynamicShapesinDeepLearningGraphsAlexanderCollinsNVIDIAacollins@nvidia.comVinodGroverNVIDIAvgrover@nvidia.comAbstractAxonisalanguagethatenablesshapeandrankinferencefortensorsinaDeepLearninggraphs.Itaimstomakeshapesimplicitandinferred,inasimilarmanner...

展开>> 收起<<

Axon A Language for Dynamic Shapes in Deep Learning Graphs_2.pdf

共9页,预览2页

还剩页未读，继续阅读

声明：本站为文档C2C交易模式，即用户上传的文档直接被用户下载，本站只是中间服务平台，本站所有文档下载所得的收益归上传人(含作者)所有。玖贝云文库仅提供信息存储空间，仅对用户上传内容的表现方式做保护处理，对上载内容本身不做任何修改或编辑。若文档所含内容侵犯了您的版权或隐私，请立即通知玖贝云文库，我们立即给予删除！

Axon A Language for Dynamic Shapes in Deep Learning Graphs_2

相关推荐

开通VIP享超值会员特权

作者详情

相关内容

热门标签

举报选择: