Named Tensor Notation
Authors: David Chiang
Alexander M Rush
Boaz Barak
Venue: Accepted by TMLR ()

Abstract: We propose a notation for tensors with named axes, which relieves the author, reader, and future implementers of machine learning models from the burden of keeping track of the order of axes and the purpose of each. The notation makes it easy to lift operations on low-order tensors to higher order ones, for example, from images to minibatches of images, or from an attention mechanism to multiple attention heads.

After a brief overview and formal definition of the notation, we illustrate it through several examples from modern machine learning, from building blocks like attention and convolution to full models like Transformers and LeNet. We then discuss differential calculus in our notation and compare with some alternative notations. Our proposals build on ideas from many previous papers and software libraries. We hope that this document will encourage more authors to use named tensors, resulting in clearer papers and more precise implementations.

Published in Transactions on Machine Learning Research (1/2023)
Named Tensor Notation
David Chiang University of Notre Dame
Alexander M. Rush
Cornell University
Boaz Barak
Harvard University
Reviewed on OpenReview:
Abstract
We propose a notation for tensors with named axes, which relieves the author, reader, and future implementers of machine learning models from the burden of keeping track of the order of axes and the purpose of each. The notation makes easy to lift operations on low-order tensors to higher order ones, for example, from images to minibatches of images, or from attention mechanism to multiple attention heads.
After brief overview and formal definition of the notation, we illustrate it through several examples from modern machine learning, from building blocks like attention and convolution to full models like Transformers and LeNet. We then discuss differential calculus in our notation and compare with some alternative notations. Our proposals build on ideas from many previous papers and softwarc librarics. We hope that our notation will encourage more authors to use named lensors, resulting in clearer papers and more precise implementations.
1 Introduction
Formal descriptions of neural networks primarily adopt the notation of vectors and matrices from applied linear algebra (Goodfellow et al., 2016). When used to describe vector spaces, this notation is both concise and unambiguous. However, when applied to neural networks, these properties are lost. Consider the equation for attention as notated in the Transformer paper (Vaswani et al., 2017):
QK VTz
Autention(Q, K) softmax
V.
The equation relates Q, K, and (for query, key, and value, respectively) as sequences of feature vectors, packed into possibly identically-sized matrices. While concise, this equation is ambiguous. Does the product QK sum over the sequence, over the features? We know that it sums over columns, but there is not enough information to know what the columns represent. Is the softmax taken over the query sequence or the key sequence? The usual notation does not offer an answer. Perniciously, the implementation of an incorrect interpretation might still run without crrors. With the addition of more axes, like multiple attention hcads or multiple sentences a minibatch, the notation becomes even more cumbersome.
We propose an alternative mathematical notation for tensors with named ares. The notation has a formal underpinning, but is hopefully intuitive enough that machine learning researchers understand it without much elfort. In named tensor notation, the above equation becomes
Attention: R key
I
xaseykkyy
vava
We follow NumPy using the azís. Other possible terms would inder, dimension, way, mode (Tucker, 1964). but felt. that anis had the least potential for confusion.

---

Published in Transactions on Machine Learning Research (1/2023)
Q ( K key
Attention(Q,K,, = ssftmax
+ V. seq
v k/eyy|
seq
The type signature introduces three named axes: the key axis is for features of queries and keys, the val axis is for features of values, and the seq axis is for tokens in a sequence. (Please see Section 2.2 for an explanation of our naming convention.) This notation makes the types of each input tensor explicit. Tensor Q is a query vector that compared with key vectors, so it has a key axis. Tensor K is a sequence of key vectors, so it has seq and key axes. Tensor V is a sequence of value vectors, so it has seq and val axes. Unlike with matrix notation, the reader is not required to remember whether seq corresponds to rows or columns in either of these lensors.
The function itself uses the named axes to precisely apply operations. The expression Q O K is a dot product
key
over the key axis shared between K and Q; there is no ambiguity about rows or columns. Similarly, the softmax function is annotated with the axis along which it is applied, removing any ambiguity or reliance on convention.
Furthermore, named tensor notation naturally extends to lifting (also known as vectorizing and/or broad- casting) function to tensors with more axes. For example, if instead of being tensor with the single axis key, Q has three axes key, seq and batch (corresponding to tokens of sequence and examples in minibatch, respectively) then the Atlention function works as written, acting on each example in minibatch in parallel. Similarly, we can also add a heads axis to the imputs to get multiple attention heads. These additional axes are often clided in neural network papers, possibly avoiding notational complexity, but possibly also hiding critical model details.
This work proposes a mathematical notation for named tensors and a fully specified semantic interpretation for the notation. Through examples, demonstrate that this notation enables specifying machine learning models and operations in a succinet yet precise manner. The need for named tensors has been recognized by several software packages, including xarray (Hoyer & Hamman, 2017), Nexus (Chen, 2017), tsalib (Sinha, 2018), axisarrays (Bauman, 2018), Named (Rush, 2019), PyTorrh (Torch Contributors, 2019), Dex (Paszke et al., 2021), JAX (JAX authors, 2021), einops (Rogozhnikov, 2022), and torchdim (DeVito, 2023). While our notation is inspired by these efforts, our focus on mathematical notation to be usod in papers, whercas previous efforts have focused on code. Our hope is that our notation will be adopted by authors, leading to clearer, more replicable papers, and that this, in turn, will encourage more implementers to adopt named tensor libraries, leading to clearer, more correct code.
Our contributions.
2 Named Tensors
In standard notation, vector, matrix, or tensor is indexed by an integer or sequence of integers; if it has dimensions Ri,..., Mr, it can be thought of as map that takes input (i1,..., i,) E [m1] x .... [n.] and outputs real number (or an element of a different field). For example, if A E R3x3, then the order of the two axes matters: AI,3 and A3,1 are not the same element. It is up to the reader to remember what each axis of each tensor stands for. This problem is exacerbated in modern machine learning, where tensors have multiple axes with different meanings (batches, channels, etc.), and different operations act on different axes.
In contrast, we propose named lensors, in which each axis has a name that describes it and ensures there is no confusion between axes. We write ax[n] for an axis with name and size n, and we write ax(i) to index the i-th clement along axis ax. So if a tensor has axes ax ax, (with ax ax, being distinct names), it can be thought of as a map that takes as input record {ax, (i1),..., (i.)), with i E [rr],...,,, ér E [rr), and outputs a field element.
In summary the key difference is that, while a tensor in standard notation takes as input an ordered tuple of indices, named tensor takes as input record, which is an unordered set of named indices. We illustrate with some examples below, then give formal definitions.
2

---

Published in Transactions on Machine Learning Research (1/2023)
2.1 By example
For example, if A represents 3 x 3 grayscale image, can make it named tensor like so (writing it two equivalent ways to show that the order of axes does not matter):
E [h xwidth]3] [wiittfjj/xheighht///
width
height
3
4
3
1
2
height
1
5
width
5 6
(1)
2
6
5]
[
9 5
We access elements of A using named indices, whose order again does not matter: Aneight(1) widthh(3 (eitth(3),,eight(() = 4. We also allow partial indexing:
width 1
height
Aheight(1)
4
Awidth(3) [4 9 5]
It does not matter if we write Aneight(1) Awidth(3) as row and column vectors. In many contexts, axis name is used with only one size. If so, we simply write height for the unique axis with namn height, in Wheightxwidthh We can leave the size of an axis unspecified at first, and specify its size later (e.g., deferring it to an appendix on experimental details). For example, we can specify hheight] = |width = 28 we want to prescribe the precise size of an image, or just write height) = width] to specify that it's a square image.
2.2 What's in a name?
Although users of this notation are free lo choose any names for axes, we offer the following recommendations. First, we recommend words instead of single letters, to communicate better the meaning of each axis.
More subtly, we rocommend words that describe a whole rathor than its parts. For example, represent minibaich of examples, we would name the axis batch; to represent sequence of tokens, we would name the axis seq. One reason for this choice is that there are cases, like height and width, where there is name for the whole, but no unanbiguous name for the part. By contrast, in cases where there is a name for the part but not the whole, it's always possible to use the plural form of the name of the part. For example, if we wanted A to have red, green, and blue channels, we would name the axis chans.
Section 4 contains many more examples axis names.
2.3 Formal definition
We now define formally the notation we use. Definition 1 (Names, indices, and axes). An awis is a pair, written ax[/], where
ax is the name of the axis, which is simply string of letters. We write both names and variables ranging over names using sans-serif font.
is set of indices. In this paper, always of the form f1....,12 n. for some n, so we abbreviate ax((1, [I[( as ax(n]..
In many contexts, there is only one axis with name ax, and so we refer the axis simply as ax. The context always makes it clear whether ax is a name or an axis. If ax an axis, we write ind(ax) for index set, and we write |ax shorthand for [ind(ax)]-
Definition 2 (Named indices and records). If ax[[]] is an axis and E I, then a named inder is a pair, written ax()). A record is set of named indices {ax, (i1 ... (ir)), where ax ... ax, pairwise distinct names.
3

---

Published in Transactions on Machine Learning Research (1/2023)
Definition 3 (Shapes). A shope is set of axes, written ax₃ axy[4]] x.. x [x,[ where ax ax, pairwise distinct names. We write 0 for the emply shape. A shape defines set of records:
rec(ax, [41] x ... xx, (4,) {{
in E i......,i 4r..
We say two shapes S and T are compatible if whenever ax(1] E S and ax[J] E T, then I= J. We say that S and T are orthogonal if there is ax such that ax[[] S and ax(J] E T for any I, J. If t E rec T and S - T, then we write ils for the unique record in rec such that i/g St.
Definition 4 (Named tensors). Let F be field and let S be shape. Then named tensor over F with shape S is a mapping from r S to F. If X has shape S then we write shp X = S. We write the set of all named tensors with shape S as FS.
We don't make any distinction between scalar (an element of F) and named tensor with empty shape (an element of ).).
If E F$, then we access an element X by applying it to a record s E rec but we write this using the usual subscript notation: X, rather than Xs)). To avoid clutter, place of X (1). we usually write (x),(t).....xx.z... When named tensor is an expression like (X +Y, we index it by surrounding it with square brackets like this: [X + (axx())))..... (x,.)
We also allow partial indexing. If X is tensor with shape T and rec where S T. then we define X, to be the named tensor with shape T 7) such that, for any t E ree(T) S),
X.)e = Xaue-
(For the edge case T 0, our definitions for indexing and partial indexing coincide: gives scalar and the other gives tensor with empty shape, but we don't distinguish between the two.)
3 Operations
A significant benefit of named tensor notation is that it allows one to unambiguously specify operations that map tensors to tensors, and defines precisely how operations can be lifted when operation is applied to tensors with more axes than are present in its signature and how broadcasting happens when different arguments add different axes.
We start with the formal definition of named tensor operations and lifting, then show how this definition leads many common operations.
3.1 Formal definition
By (named) tensor function or (named) tensor operation, we mean not only functions from tensors to tensors, but also operators like negation (-), addition (+), and so on. We will extend the standard dunction/operatoo notation by allowing tensor operations be lifted to higher-order tensors.
Definition 5 (lifting, unary). Let f: FS + GT be a function from tensors to tensors. For any shape S' orthogonal both and T, we can define the lift fs of with the shape S' to be the map
|s'; psus' + GTus'
[x)( = f(Xx))
for all E Sus and s' E rec
Usually, we simply write ʃ instead of Ss¹ That is, for every tensor X with shape R R2S, we let S(X) = (X)..
If f is a multary function, we can lift each of its arguments to larger shapes, and we don't have to add the same axes to all the arguments; an axis present in one argument but not another broadcast from the former to the latter. We consider just the case of two arguments; three more arguments are analogous.
4

---

Published in Transactions on Machine Learning Research (1/2023)
Definition 6 (lifting, binary). Let f: F$ x GT H be binary function from tensors to tensors. For any shapes S' and T' that are compatible with each other and orthogonal to S and T, respectively, and such that 8' UT' is orthogonal to U, we can lift f to:
3s'UT' psSs' x GTT 7 HUUUSUT
(SSUT(X,))) =
for all X E psus', Y FTUT I E rec(S UT).
Again, we usually write f instead of ssurr"
In the following sections, we present some consequences of the above lifting rules. In particular, we show how they allow one to lift some common operations from operating on scalars, vectors, or matrices to operating on tensors with more axes, and how they correspond voctorizing and broadcasting (as implemented by NumPy and related packages).
3.2 Elementwise operations and broadcasting
Any function from a scalar to a scalar corresponds to a tensor function with signature F + F Hence lifting to any lensor shape, by Definition 5, corresponds elementwise application. For example, given the logistic sigmoid function,
E - R
( 1 + exp(-x)
we can lift it to tensors. If A [hoight3])xwitth]]]
xwidth[3] is the example tensor (1), then
width
1+exp( -3)
T+exp( 1)
1 1+exp( (
(A) height I+exp(-1)
T+exp( 5)
I+exp(-9)
[+exp 22
T+exp(-0) 6]
T+exp(-
Similarly for rectified linear units (relu(æ)) = (0))) negation, and so on.
Any function with signature R R, including binary operators like addition (+), can be applied two named tensors with the same shape. But if we apply binary function operator to tensors with different shapes, then, by Definition 6, broadcasting applios. For example, let
E height(33]
E [width(3]
2
= [1 width
x height 7
1].
(We write x column just to make the broadcasting easier visualize.) Then, to evaluate A + x, we effectively replace with tensor with copy of x for every index of axis width. Likewise for A y.:
width
width
3 2
1 1+2
4+2 2
3
+ 4
4+1
A + x = height
7
9 9+7
A = height
1
4
9+1
2+1
6+1
5 5+1
2
4
5 5+1
Similarly for other operations. We write elementwise multiplication (Hadamard product) 0
5

---

Published in Transactions on Machine Learning Research (1/2023)
3.3 Reductions
The same rules apply to functions from vectors lo scalars, called reductions. We specily which axis reduction applies to using subscript (equivalent to the axis argument in NumPy and dim in Py'orrhh), so that even after lifting, we know which axis to reduce. For example, we can define summation:
Σ - Raall| +
3x
Σx= Σ xx(3) iES
and use it on from Eq. (1):
width
Σ A= Σ
Aneight()) = 13+1+2 1 1+5+0 + 6 4+9+5
height
height
ΣA= A Σ Awidth()
[3 13+1+4 1+5+9 6+6+55
We can also write multiple names to sum over multiple axes:
A = +3+1+4+12 -3+1-4+145+92+66 Σ A= -1
height
width
But a summation with an index variable (like i or above) is standard summation over values of that variable, and a summation with no subscript is a standard summation over a set.
Other examples of reductions include:
norm X
X2
normp X
Σ XP
ax
2x
ax
2x
min X = mint Xax(i) liE 13
max X = max o X
Xax(i) ax((i)
iiE
ax
ax
mean
Σx x
var X
1
Σ
mean X)²
ax
ax
ax
ax
ax
ax
3.4
Contraction
The vector dot dot-product is function from two vectors to scalar. We write it follows:
-: paxl1] x
paxl1]
XOY Xx Yax(s)
3x
iET
When lifted to higher-order tensors, the dot-product generalizes to the ubiquitous contraction operator, which can also be thought of as elementwise multiplication followed by summation over one axis, that is,
XoY= Exor.
(2)
3x
ax
For example,
width
A ④ - Σ AO 3-221 121.7722 7+2 1.2 7 4 4.2+9.745
height
height

---

Published in Transactions on Machine Learning Research (1/2023)
3-1+1.4+44
A ② y=
Σ A
height
1-1+5-4+9-1
width
2 2.1+6.4+5.1
Again, we can write multiple names to contract multiple axes at once:
A
A= Σ A 11.1+44 -- 6-5
5+9
height
height
width
A O with no axis name under contracts zero axes and is equivalent to elementwise multiplication, which is the reason we use the same symbol for elementwise multiplication and contraction. The contraction operator be used for many multiplication-like operations: inner product
Σ xOx
T
x
height
height
width
2 1
2 4
2 1
x
④
height
7-1 1
7.4
7-1
outer product
1-1
1-4
1-1
④ width
=
Σ AQU
matrix-vector product
width
C height
A
Σ
+ A
vector-matrix product
height
A
④ width
B
AO B
matrix-matrix product (B widthh
width
A contraction of three more tensors can be written a sum. For example, the three-way inner product of vectors F, y, 2 wiidt can be written Σ 2O width
Like the dot-product from which it lifted, but unlike matrix multiplication, the contraction operator is commutative, but not associative. However, contraction does obey the following associative-like law.
X
+ SUT u
(Y 7) = (X o ④
2 Tuu
if S n shp Z = n X 0.
(3)
The special case
S shp Z #4 X will be useful in Section 5 for moving Z from inside one more sets of parentheses to the outside.
X OY (X
(4)
3.5 Vectors vectors and beyond
Functions from vectors to vectors (Rx(11] >> Rax(!l) lift to functions tensors that operate along axis, but leave the lensor shape unchanged. Such functions particularly problematic in slandard notation, which does not provide any way (to our knowledge) of specifying which axis the operation should be performed over. Such functions include:
softmax X exp X
(5a)
Σ exp X
ax
argmax X lim softmax a. 000000
(5b)

---

Published in Transactions on Machine Learning Research (1/2023)
argmin X
lim softmax ax
(5c)
For example, we clearly distinguish between two ways of performing softmax A:
width
exp 3
exp
exp 4
exp [+exp2 exp
exp 5+eexp
exp 9+exp5 exp
softmax A height height
exp 5
exp exp [xxxp exp 3+p Tooxp
exp Fexp 5-texp6
exp 4x 0x exp exp 4-texpp cxpp
exp
1+exp S-repp
width exp 3+exp ++exp exp
exp
dxat:t 3-4
6+ Ltxxp
dho
softmax A height
exp
width
exp
exp2
exp ex e
exp
dxa x o
exp
24 exxpp 6
exp 2 6
exp 2+ 6+
3.6 Renaming and reshaping
It's often useful to rename an axis (analogous to a transpose operation in standard notation). We can think of this as the lift of a function from vectors vectors, but with different input and output axes:
[ Dx-xxxx |ax[
gax
[ +ax' (xx
X ax())
For example,
width
3
1
4
Aneight-sheight height'
1
5
2
6
5]
We can also define notation for reshaping two or more axes into one axis:
layer
Authors who need particular ordering may write a Ahheight.width) llayer
4
The order of elements in the new axis is undefined. more specific definition.
3.7 Indexing?
NumPy and its derivatives provide various ways recombine elements of a tensor form new tensor: integer array indexing, and functions like numpy take, numpy take along. axis, torch. index select, and torch gather. Using named tensors, write nearly all of these operations as lifts of a single function:
index: TR [ x [n] -> IR.
ax
(ndex(X,i) Xxx(i)-
ax
For example, suppose we have
E E @oocabfn]xemt
E [n]
1 E [n]]eq
P E pseqxvocabfu)
We grateful Tongfel Chen and Chu-Cheng Lin for contributing original idea behind this section, as well as the example.
8

---

Published in Transactions on Machine Learning Research (1/2023)
Tensor E contains word embeddings for all the words in the vocabulary. Integer is the numeric identifier of word, while tensor is a sequence of numeric identifiers of words. Tensor P contains sequence of probability distributions over the vocabulary (e.g., the predictions of language model). Then:
index(E,i) broadcasts the emb axis from to i, giving the word embedding of word i. This the vocab
same as partial indexing (Fvocab(t))-
index(EE1) also broadcasts the seq axis from / to E, giving sequence of word embeddings. This vocab is the same integer array indexing (E(TT), numpy. take (E, I, 0), or torch. indeiindeeiile sedex__elect((F, 0, !).
(ndex(P.I) aligns P's and I's seq axes, giving sequence of probabilities. This is the same vocab
numpy. take_ (P, I, 0) torch, gather (P,, 0, I).
If P and I additionally had batch axis (before the other axes), then index(P,I) would be the same vocab
wwtssorfloo.ggthhee (P, I, axis=1, batch sitch_dims=1)..
In NumPy, indexing using two or more integer arrays requires a. special definition with some surprising special cases. With named tensors, we simply apply the indexing function twice. For example, we wanted to get probabilities of words at subset. I of positions, we could let:
seq = m
E [m] subseq
positions
J E [x]] jubseq
numeric identifiers
S = index (iddxxindex(P..)..)) R subses
vocab
seq
so that
(kb)seq()) = P ses (ntteent )ocab
4 Worked Examples: Neural Networks
In this section we give a series of worked examples illustrating how standard neural network model com- ponents can be written using named tensors. Appendix builds some of these components into complete specifications of the Transformer and LeNet.
4.1 Feedforward neural networks
A multi-layer, feedforward neural network with different-sized layers be written
0 E input
X = (w
X + bb)
W R hidden,, Xinput
b2 E I hidden;
input
X² = a o
i hidden,
x + b2)
W hhidden, xhidden,
62 E IR hidden
X (W''
X x + b²)
3 goutxhideen,
3 Iout
hiddeng
The layer sizes can be specified by writing hidden, n1, etc. As noted above, a is applied elementwise and does not require additional annotation.
9

---

Published in Transactions on Machine Learning Research (1/2023)
Alternatively, the layer equation can be abstracted by defining: @
FullConn' (a) = a W
(W
x + 16
layer
layer' Soyer--layee layer
where
W E [layer/jnjxxaayrrfmm-.
b' E R glayer glayer[[d]] [nr]
The function FullConn' encapsulates both the equation for layer as well as its parameters W (analogous to what TensorFlow and PyTTrrc call modules). Since we chose to use the same axis name layer for all the layers (with different sizes ni), FullConn's temporarily computes its output over axis layer', then renames it back to layer. The network can be defined like this:
X glaver[no]
x FullConn' (X²)
X² FullConn (X²)
X FullConn (XX))
4.2 Recurrent neural networks
As a second example, we consider a simple (Elman) RNN. This model is similar to the feedforward network, except that the number of timesteps variable and parameters are shared over time. At each time step, it produces tensor with a new axis hidden' which is then renamed hidden for the next step in the recursion.
x' E R input
t = 1....,
Whe I hhidden hidden
[hidden| |hidden'
W
Innputxchiddnn
b E hhhddden'
Ro I hidden
h' W 1-11+
+
Wl x² 4b) hidden'=hhidden
1 1, I,...,m
hidden
input
4.3 Attention
In the introduction (81), we described difficulties in interpreting the equation for attention as used with Transformers (Vaswani et al., 2017). In our notation, it looks like this:
Attention:
Rley seqkey vva val
(6)
Q K key
Attention(Q,Ke (,) softmax
C V.
(7)
[key]
This definition takes a single query Q vector and returns single result vector (and actually could be further reduced to a scalar values val is not strictly necessary). To apply to sequence, we give Q seq' axis, and the function will compute output sequence. Providing Q, K, and V with a heads axis lifts the function to compute multiple attention heads.
For Transformers we often need to apply a mask to ensure altention only applied to valid keys (e.g. for causal language models). We can modify the equation to include this mask:
Attention:
keyy x keg x seqxval
seeq vaal
10

---

Published in Transactions on Machine Learning Research (1/2023)
о K
key key
(, softmax seq
+
M
@ V.
seq
Appendix A.1 includes full specification of the complete Transformer model using the named tensor nota- tion.
4.4 Convolution
Standard neural network convolutions can be specified by "unrolling" vector and then applying a standard dot product. We define axis-pa axis-parameterized unrolling function that converts one-axis tensor to sequence of kernel sized vectors:
unroll: I seq[fn] + [ [kernel] kernel
seq
kernel
unroll X Y, where
seq kernel
Y
kermel)))))))) sienii+---)
A Id convolution with input channels chans and output channels chans' consists of unrolling along the seq axis and then taking dot product:
Convid: [chansssog[x]] channs xseqin')
Convld(X; W,b) W
unroll X +
chans
where
W E Rchann xchansskeene
b E gecann'
Unrolling easily generalizes to higher-dinennsonall convolutions:
Conv2d: gechansxhiightt (h] [witth/w]] >> chans' xheight[l' xwidth iw"]
Conv2d(X; =
3
unroll unroll X +
height
kh.lew
kh
kw
where
W E Iichan'' xchansxkhxkkw
pechans
4.5 Pooling
Pooling is similar to convolutions.
We first define function to partition tensor into windows.
pool kernel
1 [eq[[]] + |gsealn/ [kkmell],
pool X Y, where
seq, kernel
Yeeg(() (olq(i))kerellll kernel(j) Xseq(ti- [kernell-j)
11

---

Published in Transactions on Machine Learning Research (1/2023)
Then we can define aggregations over kernel. We define max-pooling as: MaxPoolldæ: Rem][] + [
(X) = max pool X
kernel seq, kernel
kernel k
MaxPool2des..e
[height(A) xwidth/ui [w]
i (ee /kh) xwidthi (we/kw]
Max
max pool
pool X
kh, height, width,
|kh| kh
kw]
kw
4.6 Normalization layers
Normalization layers are used in all large-scale deep learning models, with different. architectures requiring different types of normalization. However, despite their importance, the differences between them are often not clearly communicated. For example, the PyTorrc documentation (Py Torch Contributors, 2022) describes all of them using the same equation (where > 0 is small constant for numerical stability):
Y X mean (X)
vva(X) + +
Wu & He (2018) give essentially the same equation and explain the difforences using a combination of equations, words, and pictures. But they do not capture differences in y and 3 among different normalization layers.
Critically, the layers do differ by which axes standardized well as their parameters. We define single named standardization function as:
standardize: IRax - Rax
aK
X mean X)
standardize
ax
2x
var X) +
Then, can define the three kinds of normalization layers, all with type jgbatch jaatch x llaas While superficially similar, these functions differ in their standardized axes and their parameter shape.
PatchNorm((x7y. =
standardize( (X) 10
7.8 E chans
batch,layer
InstanceNorm(X:17, standardize( @
7/ 8 chann
layer
LayerNorm((:77.9) = stlandadiiee())00
7, B E IR laane
chans
5 Differential Calculus
In many machine learning applications, need to compute derivatives of functions from tensors to tensors. In standard vector/ cattor/matree matrix notation, this can become complicated. For example, if f maps from vectors to vectors, then the partial derivatives of ʃ form a matrix (the Jacobian). It has an "input" axis for the directions in which X could change, and an "output" axis for the directions in which f(X) could change. But there are conflicting conventions about whether the first axis is the input axis ("denominator layout") or the output axis ("mumerator layout"). The derivative of a function from vectors to matrices matrices to vectors cannot. be represented matrix at all, so must resort to flattening the matrices into vectors.
12

---

Published in Transactions on Machine Learning Research (1/2023)
With non-named tensor index notation, taking derivatives is not difficult (Laue al., 2018), but again convention must be adopted that the input axes come after the output axes, separated by comma.
With named tensors, axes are not ordered, so we don't need to remember whether the input or output axes come first. But we do need to ensure that the imput and output axes have different names.
5.1 Definition
Definition 7. Let = ax, X ... x ax, be a shape. Then we write S = axi x... ... x ax,: and if s = (ax, (x,),,...xx.((. (i)) re then we write 8" = (xx())))...))((()))) Furthermore, if X RS then we write X* Xs.
Definition 8. Let f: R$ RT. The derivative of f(X) with respect to X is the tensor such that
0f(X)
0x E RRSTTT
[x((x)
0f(X),
ax* 5", 0X.
for all 8 E r and E rec Tec
The above definition assumes that S* and T are orthogonal; if not, the axes in S should be renamed to something else. For example, the second derivative (the Hessian) could be
f())
0X* oxi 8 SXX) 0x, 0x,
(XX(
ax* oxt
for all r, S and t E T.
5.2 Differentials
We could derive rules like the chain rule and the sum and product rules, and use them to compute derivatives; however, ensuring that input and output shapes are orthogonal is inconvenient. Instead, we recommend the method of differentials (Magnus & Neudecker, 1985), which reduces renaming to a minimum.
The first-order Taylor approximation of f around X R is
0f(X) O H
s(X + ~ f(X)
HERS
0x
S'
The differential of f(X) with increment H, written 0f(X; H), is the second term of this approximation; defer formal definition to Appendix B.
For example,
If id the identity function, then
id(X H) + H Oid(X; H) H.
(8a)
If f(X) = X where X E Rxx then
ax
(X H)
= = X +2X 3 H + H
ax
8x
8x
0f(X;H) 2X H.
(8b)
13

---

Published in Transactions on Machine Learning Research (1/2023)
It's often more convenient to work directly with the expression X X instead of f(X), and to write ((Xox)
for 0f(X; H). Then, since ax Ōid(X;; H1) H,, we can write Eq. (8b) simply as
a(x @ X)= = 2X xxx
ax
ax
that the H has beeen "hidden" inside 0x. More generally,
derive rules like the following:
a(u + V) = AU + avv
(9a)
a(U ④ V) = U Ov + +V o aU
(9b)
V OU U ov
0
(9c)
V2
Σ U Σ aU
(9d)
a @ V) = U O OV + V ④ 0U
(9e)
ak
aU, [OU
(9f)
OU ax- axx
[0U
(9g)
3x-+ax
The chain rule for differentials is
(x(x) ② OUs
aS(U) ax*
J: R - RT.
(9h)
X=U
Recall that can be lifted to shapes larger than S. In that case, the rule above still applies, but note that the contraction will still be over S. A special case of this is when S = =
df(x)
as(U)
OU
R:R-R.
(9i)
dr
la=t
For example, letting f(x) = exp x gives the rule
a(exp U) exp U ① aU.
(9j)
Using these rules we can compute the differential of a wide variety of expressions. For example, the softmax operator:
X) (ba)
exp X
O) (softmax
a
Σ exp X ax
E exp X)
Mexp X) exp X 3[ exp X)
(9c)
Eeep
Σ exp X) @ (exp X) exp X C Σ Xexp (9d)
E eEexpX))
> ex O exp xx exp O
(exp X Ox
(9))
xx
2x
C exp X)2
xx
Σ X) exp 0x exp X ② (exp X @ 0x)
Σ exp X)2
14

---

Published in Transactions on Machine Learning Research (1/2023)
exp X Σon expx X ax
expX Σ 3x
exp X exxp X : dx
0X
Exxp X
(5a)
softmax X @ (softmax X @ 0x) softmax X 2 0X-
ax
ak
ax
ax
softmax X (0X softmax X O 0
(10)
We stop when the only differentials left are 0X.
5.3 Derivatives via differentials
If we get (f( into so-called canonical form,
0f(x) = 0X" const.
(11)
S*
where "const." stands for terms not depending on 0x, then by Magnus Neudecker's first. identification theorem (Theorem 1 in Appendix B), we can conclude that
0f(X) D.
ax
When trying get expressions into canonical form, one helpful fact that renaming can be thought of contraction with an identity matrix. First we define the identity matrix with shape ax ax':
[1 i i=j (Tax (ax(1)) =10
i i j.
Then for any tensor A with axis ax,
Axx Taxoox ① A.
(12)
More specifically, if Ox S then
0X = 1s,s- O 0x
(13)
s"
and then Eq. (4) can usually be used to move the @ A Above, we found the differential the softmax function; now let us find its derivative.
0Xs-s- to the outermost level of the expression.
S'
A(soltmax X) (10)
softmax X @ (aX- sofumax X O OX
ax
(13)
softmax X + axx ax* softmax X O (Jax, C 0x*))
ax
ax"
(4) (softmax X o (Xxx, softmax X O lax,xx C 0X*
ax
ax
ax'
(12)
(softmax X @ (Iax.ax* (softmax x)))) O 0x*
ax'
This is in canonical form, so we have
i
(softmax. X) softmax X @ (Lax,ax* (softmax X)*).
(14)
0x
15

---

Published in Transactions on Machine Learning Research (1/2023)
5.4 Lifting
Recall that 1 is the lift of f:: I - RT' with 8', and in most contexts can simply write ʃ instead of f However, derivatives are one place where some extra caution is in order. To lighten notation, let's write for the derivative of f:
9(X) 0f(x)
0x*
Recall that the chain rule (9h) works under lifting, so
( (x)= 98 (X)
gx) @ 0Xs-s
But the contraction is only over S so it would be incorrect to conclude that 25 (x)
(X). The
8x
derivative of a lift is not the lift of a derivative. We must rename and contract S' as well:
oss¹ (x) (9h)
(X) 0Xs+8'
(13)
(X) (Is S + (Ssus
s'
( (X) s
0x sX
(SuS')
0s (X) of (X) (X) Ss
0x
In general, then, the derivative of lift is the lift of the derivative, multiplied by the identity matrix for the new axes. Intuitively, this is because the derivative is a linear transfornation beeooe lifting, transformation from S to T. When lifting to S U'S this transformation must also be lifted, which is what multiplication by Sr accomplishes.
5.5 Extended example
As more elaborate example, we find the derivative of self-attention. For brevity, we omit the factor and we write a softmax(( Q K).
keyl
seq
AA 77) a(a V)
seq
(9c)
OV +
@ da.
(15)
seq
seq
Focus first on the first term, which is the only term depending on oV:
(13)
Ov
(Jseq
c Ival.val))
@ @V*)
seq
seq
seq"
val"
((a
① /soq,soq')
O Ival C OV*
soq
seq
vall
(12)
(aseq +seq"
Ival.val") O 0V*
seg"
val*
AM, K,V) assq >sog + Ivalvvll
16

---

Published in Transactions on Machine Learning Research (1/2023)
Next, focus the second term of Eq. (15):
(10)
C
da
(a
(a(Q)
C K) a
@
a(Q
C K)))
seq
soq
key
seq
key
(a ØK K
c
aQ
0K + K
@ (Q))).
(16)
seq
key
key
seq
key
key
Keeping only terms depending on aQ, we get
(a
K
@
aQ=
a KK aQ)))
seq
key
seq
key
(13)
(a
(K
(Ikey, key"
①
oQ*) - a
@
(K O
(Ikey
key
O
aQ*)) ))))
seq
key
key"
seq
key
key
(4)
V
(a ②
(K
key"
(K
key 111
aQ
seq
key
seq
key
key"
(4)
V
c (a ④
Kkey-skey
K
Ok @
-key-skey"))
2 aQ
seq
seq
key"
V
o
(a
keey
+ K Kkey
aQ
Similarly, keeping only terms from Eq. (16) depending on 0K, we get
V (a
OK
a o (Q aK)))
seq
key
seq
key
(a
Iseq -
C (Q" C
Iseq 111
o OK* key" seq"
seq
seq
aK* =
(Q, K,
(a O Q
Jsoq,5as
Beeq,saq'
3 (Q* 2
/seq.sqq ))).
seq
seq
6 Alternatives and Related Work
6.1 Index notations
Among alternatives to standard vector and matrix notation, the most common one is index notation used in physics (Rieci & Levi-Civita, 1900). Related notations are used in other fields as well (Harshman, 2001). In this notation, axes are ordered, and every equation is written in terms of tensor components. If an index appears on both sides of an equation, then the equation must hold for each value the index, and the Einstein summation convention (Einstein, 1916) is that if an index appears twice on one side and not on the other, there an implicit summation over that index.
Attention: dd x guxde Baxd + Rd
(2 Vdz
[Attention( Q, K, V)]a = softmax
Vik-
Because A appears both sides, the equation must hold over all values of this index. But because and j occur twice only the right-hand side, they are both summed over. We would have to define exactly what the i under the softmax means (i is bound inside the softmax and free outside it), and since softmax doesn't 't distribute over addition, we would need to modify the summation convention so that the sunmation over j occurs inside the softmax.
Aside from these correctable issues, this notation scales very well to more than two axes and concise and unambiguous. But it does not solve the main problem we out solve, which is that ordered axes force the author and reader to remember the purpose of each axis. The indices do act symbolic names for axes
17

---

Published in Transactions on Machine Learning Research (1/2023)
(indeed, in abstract index notation (Penrose & Rindler, 1984), they really are symbols, not variables), but they are temporary names; they could be totally different in the next equation. It would be up to the author to choose to use consistent names, and to do so correctly.
A second issue is that because it depends on repetition of indices to work, index notation can be more verbose than our notation, particularly for reductions and contractions:
C max A,
C max A ax C = A @ B. ax
C = A,B;
Finally, index notation requires us to write out all indices explicitly. So if we wanted to lift attention to minibatches (b E [B),, multiple heads (h [H] and multiple query tokens (i' E [n'I), we would write:
Attention: BB xdv x R BBxHxxxd, x B BBxxxxxxd B
@ohnr Kohij ddl
[Attention(Q] K, V) softmax
Volak.
We could adopt a convention that lifts a function on tensors to tensors that have extra axes to the left, but such conventions tend to lead to messy reordering and squeezing/ squeezing/ussquezinn unsqueezing of axes. Named axes make such conventions unnecessary.
6.2 Graphical notations
In the graphical notation of Penrose (1971), node in the graph stands for a tensor, and its incident edges stand for its indices. The edges ordered from left to right. An edge connecting two nodes denotes contraction. The notation of Alsberg (1997) is similar, except that edges named, not ordered.
Graphs are commonly used in machine learning for representing probability models (Koller & Friedman, 2009). A node in the graph stands for a random variable, and an edge or hyperedge stands for a dependency between variables. If random variables have finite domains, then a (hyper)edge with r endpoints can thought of an r-th order tensor. A graph can then be thought of as a product and contraction. Extensions (hat allow for choice between two subgraphs (e.g., Minka & Winn, 2008) can be thought of addition.
Our assessment of graphical notations like these is that, the positive side, they have obvious value for visualization, and they at least have the potential to represent indices in purcly unordered way. On the negative side, these notations seem best suited for representing linear functions, and even for this purpose, some other practical considerations are that drawing pictures requires more effort from the author, and that pictures will have a less transparent relationship with their implementation in most programming languages.
6.3 Relational algebra
In relational algebra (Codd, 1970), the basic objects are sets of r-tuples, which could be thought of as tensors of order with Boolean-valued entries. In the original formulation, the members of the tuples, which correspond to axes, were both named and ordered, although later definitions (e.g. Pirotte, 1982) made them unordered.
Probabilistic variants of relational algebra also exist (e.g. Dey & Sarkar, 1996; Fuhr & Rölleke, 1997), whose relations would correspond to tensors of probabilities.
While relational algebra and tensor notations are designed for totally different purposes, the notation of relational algebra generally has a similar flavor to ours (for example, our contraction operator is similar the da operator, and renaming operator is the same the p operator).
18

---

Published in Transactions on Machine Learning Research (1/2023)
6.4 Programming languages
One response to the notation presented here, well the alternative notations mentioned in this section, is that research papers in machine learning should simply present models code rather than equations. But we argue that model's mathematical specification should abstract away from details of its implementation.
Conceptually, it is important to have distinet specification to define what makes an implementation (both the original implementation and any reimplementations) correct or incorrect. If the implementation is its own specification, it cannot be correct or incorroct; it will be "not even wrong."
Practically, abstracting away from implementation is important because do not want the interpretation of research papers to be subject to differences across programming languages and libraries, versions thereof. For example, PyTorch's 's Dropout2d on order-3 tensors has one behavior in versions 1.11 and 1.13, but another behavior in 1.10, 1.12, and future versions. It would be problematic for correct understanding of paper to depend on such differences.
7 Conclusions
Named tensor notation is a system of formal notation for representing operations between tensors a non-ambiguous way while remaining intuitive for practitioners. The system is motivated by challenges that arise from taking notation designed for applied linear algebra and using it for representing neural networks, demonstrated through examples of canonical deep-learning components such attention and layer normalization. However, named tensors are not limited to specifying neural networks. We have also explained how to integrate our notation with Magnus & Neudecker (1985)'s 's method of differentials for matrix calculus. While there are other conventions that such as index notation that have some usage in the machine learning community, these conventions either lack the conciseness of named tensors or are not well-suited to non-lincar operations. For these reasons, we encourage members of the machine lcarning community try out named lensor notation for leaching, research, and software documentation.
Acknowledgements
We would like to thank Ekin Akyürek, Justin Bayer, Tongfei Chen, Chu-Cheng Lin, Colin McDonald, Adam Poliak, Matt Post, Chung-chieh Shan, Nishant Sinha, and Yee Whye Teh for their input to this document (or the ideas in it). We also thank the anonymous TMLR reviewers for their feedback, which substantially improved the quality of the paper, especially Section 5.
This material based upon work supported by the National Science Foundation under Grants No. CCF- 2019291 and 213-213, as well as Sloan Fellowship, Simons Investigator Fellowship, DARPA grant 2911NF2010021, and DOE grant DE-SC0022199. Any opinions, findings, and conclusions recommen- dations expressed in this material are those of the authors and do not necessarily reflect the views of the funding agencies.
References
Bjørn K. Alsberg. A diagram notation for n-mode array equations. Journal of Chemometrics, 11:251 266, 1997. doi:
Matt Bauman. Axisarrays, 2018. URL http:://gghuu.ocoooo///aaara//rs j1. Open-source software.
Tongfei Chen. Typesafe abstractions for tensor operations. In Proc. 8th ACM SIGPLAN International Symposium on Scala, pp. 45 50, 2017. doi:
E. Codd. A relational model of data for large shared data banks. Communications of the ACM, 13(6): 977-387, June 1970. doi: 10.1145/362250 14444365334433628555
19

---

Published in Transactions on Machine Learning Research (1/2023)
Zachary DeVito. Named tensors using first-class dimensions in PyTorch, 2023. URL https://gittub..c. com/ facebokressarcchtoorchhe Open-source software.
Debabrata Dey and Sumit Sarkar. A probabilistic relational model and algebra. ACM Transactions on Database Systems, 21(3):339 369, September 1996. doi:
A. Einstein. Die Grundlage der allgemeinen Relativitatssheorie. Annalen der Physik, 354(7):769 880, 1916.
doi: 10.1002/andpp http:///nnd..11.35/4...
Norbert Fuhr and Thomas Rölleke. A probabilistic relational algebra for the integration of information retrieval and database systems. ACM Transactions on Information Systems, 15(1):32 66, Jamuary 1997. doi: 10.1145/ 10.1145/239044..2 239041.
lan Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016.
Richard A. Harshman. An index formalism that generalizes the capabilities of matrix notation and algebra to n-way arrays. Journal of Chemometrics, 15:689-714, 2001. doi: 10.1002/cem.666
Stephan Hoyer and Joe Hamman. xarray: N-D labeled arrays and datasets Python. Journal of Open Research Software, 5(1):10, 2017. doi: 10.5334/jors.144
JAX authors. Named axes and easy-to-revise parallelism, 2021. URL http:s//ja..eeddthhdess.. httpe////ttttoooo////ttttttttttt httm..aaraall...................
Daphne Koller and Nir Friedman. Probabilistic Graphical Models: Principles and Techniques. MIT Press, 2009.
Soeren Laue, Matthias Mitterreiter, and Joachim Giesen. Computing higher order derivatives of matrix and tensor expressions. In Proc. NearIPS, pp. 2750 2759, 2018. URL https https://proocediings. neuriss.eerupucce..c cc/ pdf..
Jan R. Magnus and H. Neudecker. Matrix differential calculus with applications to simple, Hadamard, and Kronecker products. Journal of Mathemalical Psychology, 29(4):474 492, 1985. doi: 10.1016/0022-
2496685))0006-9.
Tom Minka and John Winn. Gates, In Proc. NeurIPS, pp. 1073 1088-100,,000 2008. URL https::/papers..niiiii nips. cc/paper/3379-gatee. gates.
Adam Paszke, Daniel D. Johnson, David Duvenaud, Dimitrios Vytiniotis, Alexey Radul, Matthew J. Johnson, Jonathan Ragan-Kelley, and Dougal Maclaurin. Getting to the point: Index sets and parralelisn-preserviii autodiff for pointful array programming. Proc. ACM on Programming Languoges, 5(ICFP), August 2021.
doi: 10.1 10.1145; 10.1145/347733 /3473333.
R. Penrose and W. Rindler. Spinors and space-time, volume 1. Cambridge University Press, 1984.
Roger Penrose. Applications of negative dimensional tensors. In D. J. A. Welsh (ed.), Combinatorial Mathematics and its Applications, pp. 221-244. Academic Press, 1971.
Alain Pirotte. precise definition of basic relational notations and of the relational algebra. ACM SIGMOD Record, (3(1):30-45 1982. doi: 10. 10.1145/984514.98555
PyTorch Contributors. PyTorch documentation, 2022. URL http:://pytordh.orr/oos//..11 1.12.
version
M. M. G. Ricci and T. devi-Civita. Méthodes de caleul différentiel absolu et leurs applications. Mathematische Annalen, 34:125-200 1900. doi: 10.1007/BF01454220
Alex Rogozhnikov. Einops: Clear and reliable tensor manipulations with Einstein-liks notation. In Proc. ICLR, 2022. URL. https://openrevier. 7tdap/sSs
20

---

Published in Transactions on Machine Learning Research (1/2023)
Alexander Rush. Named tensors, 2019. URL http::/:.gggggha...o/////////tta
software.
Nishant Sinha. Tensor shape (annotation) library, 2018. URL Uff.:ttttee//ghhh.b..cc////ooo/t Open-
source software.
Torch Contributors. Named tensors, 2019. URL http:///pwg./.ooooooc//////////a .hmml.
PyTorch documentation.
L. R. Tucker. The exlension of factor analysis three-dimennionall matrices. H. Gulliksen and N. Fred- eriksen (ods.), Contributions Mathematical Psychology, pp. 110-127. Holt, Rinchart and Winston, New York, 1964.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Proc. NeurIPS, pp. 5998 6008, 2017. URL https: / //proceedinga iapboceednngg..eurips.....
Yuxin Wu and Kaiming He. Group normalization. In Proc. ECCV, 2018. URL https://openaccees.. thecvf. wttt://c.cccccoccceeeeeac /tt.ee//////pppee///eeaeaae Wu Noroup.oommllizaiion ECCV _cC___01V_______ Bacr8__pperpppppp
A Extended Examples
A.1 Transformer
We define a Transformer used autoregressively as a language model. The input is sequence of one-hot vectors, from which we compute word embeddings and positional encodings:
E [0,1]sexxw))
1=1
vocab
W
E C layer
E E vocabbllaye
vocab
E R seq x layer
sin((p (1/11000 odd cos((p (1- even.
Peqq(p),layer(?
Then we use L layers of sell-attention and feed-forward neural networks:
x W +
T LayerNorm's (SelfAtt' (x*) + x²)
X LayerNorm 1' (FFN" +T!)
T LayerNorm (SelfAtt (x (X--1) XL-1)
X LayerNorm' 1' (FFN' T
O soft x
vocab
layer
where LayerNorm, SelfAtt and FFN are defined below.
Layer normalization (I 1..,1'...,,L.II"" L'D:):
LayerNorm': Rlayer - Jalayer
LayerNorm' standardize(X) 0
layer
21

---

Published in Transactions on Machine Learning Research (1/2023)
g B R
We defined attention in $4.3; the Transformer uses multi-head self aelf-attention, in which queries, keys, and values are all computed from the same sequence.
SelfAtt': Reeqxlayer -> seq Xlayer
SelfAtt Y
where
seql
|seq
keyl
val|
|layer|
heads
Q = p C Xxeq
Wlo E R heaads x x
layer
K W X
Wik
gheadsxlayerxkke
layer
V
W + X
WLV
E
R
Thheads layer layer
layer
M E seeqseee
0
i
(e)s(())
DC
otherwise
Y w1o -1,
Attention(Q,Ae M seq' -tes -seq
17.0 gheadssvalllayer
heads
Feedforward neural networks:
FFN:: laleyr R glaver FXN'(X)= = X2
where
X relu(W!!!
3 X + +b
W
6 R hidden
layer
x = W42
x + 6-2
W4.2 E Iglayexchidden
6.2 R hidden
hidden
A.2 LeNet
X E R gbatchhxchansfen (aannnxcccca(ooxxeeeeeeeee x
T' relu(Conv' ))
X1 MaxPool (T*)
22 velu(Con?2(x1))
X² MaxPool (T (theigh..withh.hhhnss--t. Hlaye
3 3elu(W'
x + b²)
W hhidden
b I hidden
layer
0 softmax (W4
C hidden
x + 3
W hildsessxhidden
6 E d classes
classes
As an alternative the flattening operation in the equation for X2, we could have written
2 MaxPool2(TS)
22

---

Published in Transactions on Machine Learning Research (1/2023)
X @ x b9)
W IR hgdidden chans
height width chans
The convolution and pooling operations are defined as follows:
Conv' = Conv . Dchhanssssss channs
where
W IRchhans (chnns//s//tt_11 (xw.k xkw(/ 8 E Rhchnns [cc]
and
MaxPool MaxPool2d pt
ph .ph! (X).
B
Differentiation: Formal Definitions
The following definition and theorem come directly from the paper by Magnus & Neudecker (1985), but generalized to named tensors.
For any X E RS, we write ||x||| = norm X.
Definition Let f: S + RT where S RRS. Let A be an interior point of R that is, for some r > 0, B(Arr) [X || All < <r C S. If there is lensor D(A) RS and R(A, H) IR? such that
HA++))= f(A) + D(A) @ Hs--s- + R(A,H)
for all H E TRS with ||]]| <r, r, and
R(A, H))
lim #-0
O,
||
then f is said to be differentiable at A; the tensor
0f(A: = D(A) @ Hs-
is then called the (forst) differentiol of f at A with increment H.
Magnus& & Neudecker give their (first) identification theorem twice, once for vector-to-vector functions and once for matrix-tar to- to-matrix functions (but omitting vector-to- toar-to-mann matrix and matrix-to-vest vector functions). Here, we only need one version, which works for functions from tensors to tensors of any shape.
Theorem 1. Let f: S - TT, where S C RS, be differentiable at A S. Let D(X) E RUSTT Then
0f(X)
for all H, 0)(A: D(X) @ Hs--5-
iff
D(X).
0X*
X==A
C BTęX Macros
Many of the ыЛTХ macros used in this document are available in the style file namedtensor.es .str,, available on CTAN at https: / /ttn www.poggg/amettennniigg To use it, put
(ssepackagennamddteessooo
23

---

Published in Transactions on Machine Learning Research (1/2023)
in the preamble of your ETEX source file (after /dccumentcclasssaariiill but before \begin(doccment))) We write axis names in sans-serif font. To make this easier, \name(ax} prints an axis name (like this: ax). and Andef (aax) {ax} defines a macro \ax that does the same.
Binary operators
Use A \ndot([ax) B for contraction: A @ B. You use V\ to stack up several names.
In general, you can use Anbin to make new binary operator with name under it:
A (Vstn)) gives you A * B.
xx
Functions
Use \axum[\ax} for summation: Σ A.
In general, you Anfun to make function with name under it: (axim())xXXxxx)) A gives you qux A.
24


Reviews:
[{"review_id": "JfQ1vw53RU", "reviewer": "Reviewer_V2jH", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "* (+) the proposed notation is indeed suitable to resolve ambiguities in defining common mechanism.\n* (+) the breadth of worked examples shows that the proposed notation is suitable to handle a wide array of use cases.\n* (-) there is little novelty here, given that it summarises a trend of (referenced) similar proposals.\n* (-) the paper is focused only on notation. However, one reason for choosing to represent operations as (somewhat ambigious and hence error-prone) matrix multiplication is that this closely corresponds to factors influencing hardware performance. Keeping these factors in mind is crucial to enable scale-out and hence important in communicating results.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper formally introduces the notion of naming tensor axes/dimensions/... and shows how to adapt common tensor operations to this convention. A number of worked examples on common building blocks (MLPs, RNNs, attention, convolutions, ..., up to differentiation) show that standard machine learning tools can be concisely expressed using this notation.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "* page 1: I found $\\mathbb{R}^{\\textsf{key}}$ (etc.) surprising, as it is not refering to the dimension of individual keys etc. This aspect becomes clearer in the following, but it would be good to call this out earlier.\n* page 2: https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html is another example of an implementation of the proposed notation (which actually discusses HW implementation issues in great detail)\n* page 4, def 6/7: I'd recommend to rename $s$ to $s'$ here, as it is an element of $\\textnormal{rec}\\mathcal{S}'$ (and not the also existent $\\textnormal{rec}\\mathcal{S}$)\n* page 5, Sect. 3.2: please recall for the reader that $A$ was defined in Sect. 2.1. Given that $A$ is also used as a variable in some definitions, it may be easier for the reader if you picked another name for this concrete matrix.\n* page 7, Sect. 4.1: \"(analogous to what TensorFlow and PyTorch modules)\" misses its verb\n", "broader_impact_concerns": "I have no concerns regarding the broader impact of this paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "3PM5digS4f", "reviewer": "Reviewer_qUcr", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The proposed notation is quite clean and minimal while also being unambiguous, and I think adopting this notation would lead to more readable ML papers.\n- The notation also appears to be very general and expressive, and easy to extend if needed.\n- The examples provided by the authors both demonstrate the value of the notation and serve as a reference for future uses of the notation.\n- Overall I think this notation will be of interest to members of the TMLR community.\n\nWeaknesses:\n\n- I found some parts of the paper confusing, as discussed in the requested changes section below.\n- The definition of differential and corresponding differentiation rules seemed imprecise and not fully specified, which might lead to difficulty using this notation to express derivative computations.\n- It's not currently obvious from the paper how to actually typeset this notation (e.g. using LaTeX)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a notation for representing functions between multidimensional arrays (tensors), for use in academic papers and presentations. The specific contributions are:\n\n- The authors define a notation for representing tensors in terms of semantically named axes, and operations that act on those tensors, motivated by drawbacks of existing notation used in papers.\n- They specify the interpretation of that notation, in particular by interpreting a tensor as a mapping from a set of named indices (a \"record\") to a field element, and describing the corresponding mappings for each of their operations.\n- They show how the notation can be used to write a variety of operations common in neural network architectures.\n- They explain how to use this notation to compute Jacobian tensors, and present a set of rules analogous to the scalar differentiation rules.\n- They give examples of how to compute Jacobian tensors for a few operations of interest.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "## Important requested changes\nThere were some things I found confusing about the current version of the paper, specifically regarding the computation of derivatives. I would like to see these addressed before recommending acceptance.\n\n### Confusing/incomplete definition of differentials (Section 5.2)\n\nI found Section 5.2 to be very unclear, specifically regarding the definition of a differential. Upon first read, I was very confused regarding where $\\Delta X$ is defined, what type of object $\\partial f(X; \\Delta X)$ is, and why it is OK to omit $\\Delta X$ in the calculations; I initially thought the differential must be a function taking $\\Delta X$ as an argument, but realized this didn't make any sense once I got to the list of rules.\n\nAfter consulting Magnus & Neudecker (1985), I think that my confusion comes from two things:\n- a missing part of the definition of differential. Magnus & Neudecker define \"the differential of f(X) at X **with increment $\\Delta X$**\", where $\\Delta X$ is also explicitly named as part of the term, and different increments produce different differentials. \n- the current wording that \"the differential ... is a tensor ... that linearly approximates f\", which makes it sound like the tensor itself is a mapping somehow (since a fixed tensor can't linearly approximate a function).\n\nIf my understanding is correct, I would recommend rewording this section to be more explicit about the role of $\\Delta X$ in the differential, and explicitly naming it as part of the definition like Magnus & Neudecker do. Perhaps something like\n\n> Intuitively, the differential of $f(X)$ at $X$ with increment $\\Delta X$, written $\\partial f(X; \\Delta X)$, is a tensor with the same shape as $f(X)$, obtained by evaluating a linear approximation of $f$ around $X$ at the position $X + \\Delta X$.\n\nor even\n\n> Intuitively, the differential of $f(X)$ at $X$ with increment $\\Delta X$, written $\\partial f(X; \\Delta X)$, is the first-order term of the Taylor series for $f(X)$ around $X$, evaluated at $X + \\Delta X$.\n\nIt might be even simpler to just write out a formal definition of the differential instead of omitting it; it looks like the definition would only be a few lines long?\n\nRelatedly, I think the sentence \"In practice, we can write $\\partial f(X; \\Delta X)$ simply as $\\partial f$\" is not quite correct, and should instead be \"as $\\partial f(X)$\", since subsequent uses of $\\partial$ always apply to terms, not functions. (I also think \"never appears in our calculations\" isn't the best justification, and it would be better to say that it's omitted because it can be inferred from context.)\n\nOne more minor point on the same topic: it seems that this section is a recommendation / example of how to use the notation to compute derivatives, so I'd suggest wording it as \"we can use the method of differentials\" or \"we propose/recommend using the method of differentials\" instead of just saying \"we use the method of differentials\". At first I thought this section might be some sort of implementation detail of some autodiff system used by the authors, which was confusing.\n\n### Definition of chain rule (Section 5.2)\n\nThe current definition of the chain rule is only well defined when $\\\\mathcal{U}$ and  $\\\\mathcal{V}$ are orthogonal, but this is not explicitly stated. Moreover, I think it's important to discuss how to do the chain rule when shapes are NOT orthogonal, since this will occur often in practice. My guess is that this will involve a rename of any conflicting axes, something like\n\n$$\n\\\\partial f(U) = \\\\left[ \\\\frac{\\\\partial f(U)_{\\\\mathcal{U} \\to \\\\mathcal{U}'}}{\\\\partial U} \\\\odot_\\\\mathcal{U} \\\\partial U \\\\right]\\_{\\\\mathcal{U}' \\to \\\\mathcal{U} }\n$$\n\nperhaps. (Alternatively the rename could be applied to the denominator, I'm not sure which is simpler.)\n\n### Additional details in example derivations (Section 5.3)\n\nI found the examples in 5.3 hard to verify. While each of the steps seemed plausible, it was difficult to determine which of the rules were being applied, and whether they were being applied correctly. As a derivation of the equations themselves, this seems OK, but as an example of how to use the notation and rules, I think it is important to actually show why the steps follow from the stated rules.\n\nOne thing I noticed was the combination of the chain rule with the quotient rule in a single step; I think it would be better to split those steps up.\n\nI also noticed that the transformation from $\\partial \\exp X$ to $\\exp X \\odot \\partial X$ (hidden in the above step) involves an elementwise variant of the chain rule that is never explicitly stated. The chain rule written in section 5.2 would suggest that the differential would be  $\\partial \\exp X = \\exp X \\odot_{ax} \\partial X$, but this doesn't work, because it reduces away `ax` incorrectly (since, as a vector-to-vector function, $\\exp$ doesn't satisfy the unstated orthogonality requirement). It seems important to either explicitly discuss the chain rule for elementwise functions, or show how it follows from the ordinary chain rule after applying all of the necessary renaming. \n\nAdditionally, at multiple places throughout the examples, the canonical form of the differential is not actually shown, despite how much 5.2 stresses the importance of the canonical form. I suggest explicitly writing out the canonical form (with a contraction instead of a sum) at all of the places where Theorem 1 is used in sections 5.3 and 5.4.\n\n---\n\n## Other suggestions\nI think the paper could be improved by adding some more discussion of these areas.\n\n### Introduction of named tensors (Section 2.0)\nI found the introduction of named tensors a bit disorienting. I think part of my confusion comes from the section being written from a perspective of named axes already existing, despite the notation not being defined yet. For instance, it currently reads \"If it has ax1[n1], . . . , axr[nr] (with ax1, . . . , axr being distinct names), it can be thought of as a map that takes as input a record {ax1[i1],...,axr[ir]}, with i1 ∈ [n1],...,ir ∈ [nr], and outputs a field element.\" However, as a reader, I don't yet know what `ax1[n1]` means, or what it means for a named tensor to have it, or what a record is, and only in the following sections do I learn that `ax1` is a variable taking the place of a name, and `n1` is supposed to denote the size of that axis.\n\nI think it would be clearer to briefly define each of the concepts (informally) before using them, and only introduce one at a time. For instance, perhaps something like\n\n> In contrast, we propose named tensor notation, in which each axis has a name that describes it and ensures there is no confusion between axes. We use `ax[n]` to refer to an axis with name `ax` and size `n`, and `ax(i)` to the `i`th element along that axis (called a named index). We can then define a named tensor as a map that takes a *record* and maps it to a field element, where a record is ...\n\nAnother confusing aspect was that the record notation is inconsistent with the later definition, it should be $\\\\{ax_1(i_1),...,ax_r(i_r)\\\\}$. (I also mention this in the \"Nits\" section below.)\n\n### Elementwise operations as a special case of ordinary broadcasting (Section 3.2)\nIt seems to me that the rules for broadcasting elementwise operations follow directly from the rules for arbitrary functions, if you view elementwise operations as operations on tensors with the empty shape (as mentioned at the end of section 2.2). This seems pretty cool, and might be worth an explicit mention.\n\n### Discussion of tensor-to-tensor operations with multiple (distinct) axes\nAll of the examples in section 3 appear to treat input axes identically, e.g. summing over two axes at once is the same as flattening a sum over all elements of a reshaped axis. However, later examples use functions (such as `unroll`) which take multiple axes as arguments and apply different operations along each. It might be worthwile to add a brief mention of this (perhaps in section 3.4), since it seems like a strength of the notation that such functions are easy to express.\n\n### Discussion of gather/scatter-type operations\n\nDoes this notation support \"scatter\" / \"gather\" operations, where one tensor is used to index into another? How would such operations be represented? It would be interesting to add a discussion of those, since those are fairly common when implementing ML algorithms.\n\n\n### Motivation for differentiation section\nSection 5 (\"Differentiation\") seems a bit out-of-place currently, as it jumps directly into definitions and rules without explaining why. Is this section supposed to be another example of why the notation is useful, by showing that it aids computations? Or are the derivative rules intended to be part of the notation? It would be useful to start this section with some sort of statement summarizing what this section does and how it relates to the remainder of the paper.\n\n### Example of applying chain rule with non-orthogonal axes\nBuilding on my earlier point about the chain rule with non-orthogonal names, I think it would be useful to see an example of a derivation of a differential that involves a non-trivial chain rule application for a function that preserves some axes in its output. For instance, a computation of something like $\\\\partial ( softmax_{seq} ( Q \\\\odot_{key} V ))$. Right now the rules don't clearly explain what to do in this situation, as far as I can tell.\n\n### Discussion of other alternative notations\n\nTwo other notation options which might be worth discussing:\n\n- A version of index notation that doesn't include the Einstein summation convention, and requires all summation to be written out (a bit less ambiguous than including Einstein summation, but more verbose)\n\n- Just writing out everything in code, e.g. in numpy/pytorch/jax\n\nNot that these are particularly compelling alternatives, but I've seen them used in practice.\n\n### LaTeX definitions for notation\n\nSince this paper is recommending a new type of notation for paper-writing specifically, it would be convenient to include a reference for how to typeset this notation in LaTeX. This might help standardize some of the details (e.g. the centering of subscripts below $\\\\odot$ and the sans-serif notation for names).\n\n---\n\n## Nits / typos / minor suggestions\nI noticed a few minor things while reading through the paper:\n\n- Named tensor software: There have been some newer software implementations of named tensors that aren't mentioned here, including [einops](https://openreview.net/pdf?id=oapKSVM2bcj), [torchdim](https://github.com/facebookresearch/torchdim) and [jax.xmap](https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html). Also, the citation for Dex could be updated to refer to the [conference version](https://dl.acm.org/doi/abs/10.1145/3473593) instead of the workshop paper.\n \n- `Section 2.0`: Inconsistent notation of introductory example. \"a map that takes as input a record $\\\\{ax_1[i_1],...,ax_r[i_r]\\\\}$\" is inconsistent with the notation defined in the next section; it should be $\\\\{ax_1(i_1),...,ax_r(i_r)\\\\}$.\n  - I found this pretty confusing at first, since it seemed like the $ax_1[i_1]$ notation was being overloaded.\n\n- `Definition 5 (Named tensor operation)`: It seems that your definition of operation is equivalent to a function, and throughout the paper you use \"function\" and \"operation\" interchangeably. I wonder if you even need to formally define \"operation\", or whether \"function\" would be sufficient?\n\n- `Definition 7 (lifting, binary)`: Slightly odd grammar; perhaps \"and **for which** S' ∪ T' is orthogonal to U\" would be better\n\n- `Section 3.2`: It would be helpful to restate the definition of $A$ before applying it, or at least refer back to where A was defined (on page 2).\n\n- `Section 3.3`:\n  - If \"norm\" is intended as a convention other authors should use, it might be worth allowing other norms than the L2 norm. (Or, if this is just an example for how to define a reduction, it might be worth stating that explicitly.)\n  - Could you include an example of contracting multiple axes at once?\n\n- `Section 3.4`: I was surprised that `argmax` was defined as a one-hot embedding, that seems unusual. Is there not a way to encode the ordinary notion of `argmax` as an integer-valued tensor?\n\n- `Section 4.1`: It's not clear what it means for `FullConn` to \"encapsulate\" the parameters if the parameters are arguments to `FullConn`. It seems clearer not to have W and b be inputs to `FullConn` if the intent is to have `FullConn` implicitly carry W and b.\n\n- `Section 4.4`: I expected the bias term for convolutions to have a `chans'` axis, is it intentional that it does not?\n\n- `Section 5.4`: Missing word(s): \"But although f' can be to X using the usual lifting rules\". f' can be what to X?\n\n- `Appendix A.1`: LayerNorm is defined in terms of XNorm, but XNorm is not defined. Perhaps this was supposed to refer to \"standardize\" from the main paper?\n\n", "broader_impact_concerns": "I do not have any broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PDNvDsJrwq", "reviewer": "Reviewer_TBKC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall, this is a well-motivated and clearly presented study with useful results for researchers in machine learning. Einstein summation notation is the prevalent tensor notation (also called index notation in the paper). This notation has the advantage of clearly expressing all multilinear operations, but it is insufficient for machine learning algorithms because it cannot handle other operations (convolution, pooling, etc.). The named tensor notation, in my opinion, generalizes the index notation and would help scientific writing in the machine learning field.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work offers a named tensor notation in which each dimension/axis of a particular tensor is given a name and the names are used to create unambiguous and easily understood tensor operations. This syntax can be used to dramatically improve the readability of machine learning models, as demonstrated by multiple instances. Detailed comparisons between this notation with other existing tensor notations, including Einstein summation notation and the graphical tensor diagrams, are also given. ", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Below are my suggestions:\n\n1. The tensor contraction examples shown in the paper all involve 2 tensors/matrices/vectors. I think it would be helpful to give an example showing how to express general tensor contractions with more than 2 tensors with the named notation. For example, the contraction in the einsum notation $d_{i,k,l} = \\sum_{j}a_{i,j,k}b_{i,j}c_{i,j,l}$ can be expressed as $d = \\sum_{j}(a \\odot b \\odot c)$ in the named notation.\n\n2. Definition of derivatives under nonorthogonal case is a bit confusing, since it's not clear to me how to use that in the chain rule. For the example in 5.3, authors show how to get $\\frac{\\partial L}{\\partial X}$ based on differential. Is there easy ways one can derive it using chain rule: $\\frac{\\partial L}{\\partial X}=\\frac{\\partial L}{\\partial Y}\\frac{\\partial Y}{\\partial X}$? It's not clear to me how to use the chain rule directly here since index renaming is needed for $\\frac{\\partial Y}{\\partial X}$. I believe more discussions here would be helpful.\n\n3. Minor: at the beginning of page 13: $\\frac{\\partial L}{\\partial X}$ instead of $\\frac{\\partial f(Y)}{\\partial X}$.\n", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewers agree that the proposed notation is well-designed and expressive. The reviewers appreciated the author responsiveness in the review process and feel that after recent updates, the paper is correct and clearly relevant to the TMLR community. Improved practices around communicating tensor notation would be valuable for the community, so we hope to see widespread adoption.
Decision: Accept (poster)