https://namedtensor.github.io/

\( \require{ams} \DeclareMathOperator*{\softmax}{softmax} \
DeclareMathOperator{\ind}{ind} \DeclareMathOperator{\rec}{rec} \
newcommand{\ensuremath}[1]{#1} \newcommand{\vdotswithin}[1]{\vdots}
\)

Named Tensor Notation

David Chiang
University of Notre Dame

Sasha Rush
Cornell University

Boaz Barak
Harvard University

Version 0.3

  * 1 Introduction
  * 2 Informal Overview
  * 3 Examples
  * 4 LaTeX Macros
  * 5 Formal Definitions
  * 6 Differentiation
  * 7 Extensions
  * 8 Alternatives
  * Acknowledgements
  * References

1 Introduction

Most papers about neural networks use the notation of vectors and
matrices from applied linear algebra. This notation is optimized for
talking about vector spaces, but becomes cumbersome when talking
about neural networks. Consider the following equation (Vaswani et
al. 2017): \[\text{Attention}(Q, K, V) = \softmax \left( \frac{QK^\
top}{\sqrt{d_k}} \right) V.\] where \(Q\), \(K\), and \(V\) (for
query, key, and value, respectively) are sequences of feature
vectors, packed into matrices. Does the product \(QK^\top\) sum over
the sequence, or over the features? It sums over columns, but there's
not enough information to know what the columns represent. Is the
softmax taken over the query sequence or the key sequence? The usual
notation doesn't even offer a way to answer this question. With
multiple attention heads or multiple sentences in a minibatch, the
notation becomes more difficult still.

Here, we propose mathematical notation for tensors with named axes.
The notation has a formal underpinning, but is hopefully intuitive
enough that machine learning researchers can understand it without
much effort.

In our notation, the above equation becomes \[\begin{aligned} \text
{Attention} \colon \mathbb{R}^{\ensuremath{\mathsf{seq'}} \times \
ensuremath{\mathsf{key}}} \times \mathbb{R}^{\ensuremath{\mathsf
{seq}} \times\ensuremath{\mathsf{key}}} \times \mathbb{R}^{\
ensuremath{\mathsf{seq}} \times\ensuremath{\mathsf{val}}} &\
rightarrow \mathbb{R}^{\ensuremath{\mathsf{seq'}} \times \ensuremath
{\mathsf{val}}} \\ \text{Attention}(Q,K,V) = \mathop{\underset{\
ensuremath{\mathsf{seq}}}{\mathrm{softmax}}} \left( \frac{Q \mathbin
{\underset{\ensuremath{\mathsf{key}}}{\odot}} K}{\sqrt{|\ensuremath{\
mathsf{key}}|}} \right) \mathbin{\underset{\ensuremath{\mathsf{seq}}}
{\odot}} V.\end{aligned}\] The tensor \(K\) has axes for the sequence
(\(\mathsf{seq}\)) and for the key features (\(\mathsf{key}\)),
instead of rows or columns, so the reader does not need to remember
which is which. The dot product \(Q \mathbin{\underset{\ensuremath{\
mathsf{key}}}{\odot}} K\) is explicitly over the \(\mathsf{key}\)
axis. The resulting tensor has a \(\mathsf{seq}\) axis for the key
sequence and a \(\mathsf{seq'}\) axis for the query sequence, and the
softmax is explicitly over \(\ensuremath{\mathsf{seq}}\), as is the
dot product with \(V\). This formula works as written if we add a \(\
ensuremath{\mathsf{heads}}\) axis for multiple attention heads, or a 
\(\ensuremath{\mathsf{batch}}\) axis for multiple sequences in a
minibatch.

Our notation is inspired by libraries for programming with
multidimensional arrays (Harris et al. 2020; Paszke et al. 2019) and
extensions that use named axes, like xarray (Hoyer and Hamman 2017),
Nexus (Chen 2017), tsalib (Sinha 2018), NamedTensor (Rush 2019),
named tensors in PyTorch (Torch Contributors 2019), and Dex 
(Maclaurin et al. 2019). However, our focus is on mathematical
notation rather than code.

The source code for this document can be found at https://github.com/
namedtensor/notation/. We invite anyone to make comments on this
proposal by submitting issues or pull requests on this repository.

2 Informal Overview

In standard notation, a vector, matrix, or tensor is indexed by an
integer or sequence of integers. If \(A \in \mathbb{R}^{3\times3}\),
then the order of the two axes matters: \(A_{1,3}\) and \(A_{3,1}\)
are not the same element. It's up to the reader to remember what each
axis of each tensor is for. We think this is a problem and propose a
solution.

2.1 Named tensors

In a named tensor, we give each axis a name. For example, if \(A\)
represents an image, we can make it a named tensor like so (writing
it two equivalent ways to show that the order of axes does not
matter): \[\begin{aligned} A &\in \mathbb{R}^{\ensuremath{\ensuremath
{\mathsf{height}}[3]} \times \ensuremath{\ensuremath{\mathsf{width}}
[3]}} = \mathbb{R}^{\ensuremath{\ensuremath{\mathsf{width}}[3]} \
times \ensuremath{\ensuremath{\mathsf{height}}[3]}} \\ A &= \
ensuremath{\mathsf{height}}\begin{array}[b]{@{}c@{}}\ensuremath{\
mathsf{width}}\\\begin{bmatrix} 3 & 1 & 4 \\ 1 & 5 & 9 \\ 2 & 6 & 5 \
end{bmatrix}\end{array} = \ensuremath{\mathsf{width}}\begin{array}[b]
{@{}c@{}}\ensuremath{\mathsf{height}}\\\begin{bmatrix} 3 & 1 & 2 \\ 1
& 5 & 6 \\ 4 & 9 & 5 \end{bmatrix}\end{array}.\end{aligned}\]

We access elements of \(A\) using named indices, whose order again
does not matter: \(A_{\ensuremath{\ensuremath{\mathsf{height}}(1)}, \
ensuremath{\ensuremath{\mathsf{width}}(3)}} = A_{\ensuremath{\
ensuremath{\mathsf{width}}(3)}, \ensuremath{\ensuremath{\mathsf
{height}}(1)}} = 4\). We also allow partial indexing: \[\begin
{aligned} A_{\ensuremath{\ensuremath{\mathsf{height}}(1)}} &= \
ensuremath{\mathsf{}}\begin{array}[b]{@{}c@{}}\ensuremath{\mathsf
{width}}\\\begin{bmatrix} 3 & 1 & 4 \end{bmatrix}\end{array} & A_{\
ensuremath{\ensuremath{\mathsf{width}}(3)}} &= \ensuremath{\mathsf{}}
\begin{array}[b]{@{}c@{}}\ensuremath{\mathsf{height}}\\\begin
{bmatrix} 4 & 9 & 5 \end{bmatrix}\end{array}.\end{aligned}\]

In many contexts, an axis name is used with only one size. If so, we
can simply write \(\ensuremath{\mathsf{height}}\) for the unique axis
with name \(\ensuremath{\mathsf{height}}\), as in \(\mathbb{R}^{\
ensuremath{\mathsf{height}} \times \ensuremath{\mathsf{width}}}\). We
can leave the size of an axis unspecified at first, and specify its
size later (like in a section on experimental details): for example, 
\(|\ensuremath{\mathsf{height}}|=|\ensuremath{\mathsf{width}}|=28\)
to specify its exact size or just \(|\ensuremath{\mathsf{height}}|=|\
ensuremath{\mathsf{width}}|\) to specify that it's a square image.

What are good choices for axis names? We recommend meaningful words
instead of single letters, and we recommend words that describe a
whole rather than its parts. For example, if we wanted \(A\) to have
red, green, and blue channels, we'd name the axis \(\mathsf{chans}\),
and if we wanted to represent a minibatch of images, we'd name the
axis \(\mathsf{batch}\). Please see SS3 for more examples.

2.2 Named tensor operations

Operations on named tensors are defined by taking a function on
low-order tensors and extending it to higher-order tensors.

2.2.1 Elementwise operations and broadcasting

Any function from a scalar to a scalar can be applied elementwise to
a named tensor, and any function from two scalars to a scalar can be
applied to two named tensors with the same shape. For example: \[\
frac1{1+\exp(-A)} = \ensuremath{\mathsf{height}}\begin{array}[b]{@{}
c@{}}\ensuremath{\mathsf{width}}\\\begin{bmatrix} \frac 1{1+\exp(-3)}
& \frac 1{1+\exp(-1)} & \frac 1{1+\exp(-4)} \\[1ex] \frac 1{1+\exp
(-1)} & \frac 1{1+\exp(-5)} & \frac 1{1+\exp(-9)} \\[1ex] \frac 1{1+\
exp(-2)} & \frac 1{1+\exp(-6)} & \frac 1{1+\exp(-5)} \end{bmatrix}\
end{array}.\]

But if we apply a binary function/operator to tensors with different
shapes, they are broadcast against each other (similarly to NumPy and
derivatives). Let \[\begin{aligned} B &\in \mathbb{R}^{\ensuremath{\
ensuremath{\mathsf{height}}[3]}} & C &\in \mathbb{R}^{\ensuremath{\
ensuremath{\mathsf{width}}[3]}} \\ B &= \ensuremath{\mathsf{height}}\
begin{array}[b]{@{}c@{}}\ensuremath{\mathsf{}}\\\begin{bmatrix} 2 \\
7 \\ 1 \end{bmatrix}\end{array} & C &= \ensuremath{\mathsf{}}\begin
{array}[b]{@{}c@{}}\ensuremath{\mathsf{width}}\\\begin{bmatrix} 1 & 4
& 1 \end{bmatrix}\end{array}.\end{aligned}\] (We write \(B\) as a
column just to make the broadcasting easier to visualize.) Then, to
evaluate \(A+B\), we effectively replace \(B\) with a new tensor \(B'
\) that contains a copy of \(B\) for every index of axis \(\
ensuremath{\mathsf{width}}\). Likewise for \(A+C\): \[\begin{aligned}
A + B &= \ensuremath{\mathsf{height}}\begin{array}[b]{@{}c@{}}\
ensuremath{\mathsf{width}}\\\begin{bmatrix} 3+2 & 1+2 & 4+2 \\ 1+7 &
5+7 & 9+7 \\ 2+1 & 6+1 & 5+1 \end{bmatrix}\end{array} & A + C &= \
ensuremath{\mathsf{height}}\begin{array}[b]{@{}c@{}}\ensuremath{\
mathsf{width}}\\\begin{bmatrix} 3+1 & 1+4 & 4+1 \\ 1+1 & 5+4 & 9+1 \\
2+1 & 6+4 & 5+1 \end{bmatrix}\end{array}.\end{aligned}\]

2.2.2 Reductions

The same broadcasting rules apply to functions from vectors to
scalars, called reductions. Unlike with functions on scalars, we
always have to specify which axis reductions apply to, using a
subscript. (This is equivalent to the axis argument in NumPy and dim
in PyTorch.)

For example, we can sum over the \(\ensuremath{\mathsf{height}}\)
axis or the \(\ensuremath{\mathsf{width}}\) axis of \(A\): \[\begin
{aligned} \sum\limits_{\ensuremath{\mathsf{height}}} A &= \sum_i A_{\
ensuremath{\ensuremath{\mathsf{height}}(i)}} = \ensuremath{\mathsf{}}
\begin{array}[b]{@{}c@{}}\ensuremath{\mathsf{width}}\\\begin{bmatrix}
3+1+2 & 1+5+6 & 4+9+5 \end{bmatrix}\end{array} \\ \sum\limits_{\
ensuremath{\mathsf{width}}} A &= \sum_j A_{\ensuremath{\ensuremath{\
mathsf{width}}(j)}} = \ensuremath{\mathsf{}}\begin{array}[b]{@{}c@{}}
\ensuremath{\mathsf{height}}\\\begin{bmatrix} 3+1+4 & 1+5+9 & 2+6+5 \
end{bmatrix}\end{array}.\end{aligned}\] See SS3.1.1 for more examples
of reductions.

We can also write multiple names to perform the reduction over
multiple axes at once. For example, \[\sum\limits_{\ensuremath{\
mathsf{height,width}}} A = \sum_i \sum_j A_{\ensuremath{\ensuremath{\
mathsf{height}}(i)},\ensuremath{\ensuremath{\mathsf{width}}(j)}} =
3+1+4+1+5+9+2+6+5.\]

The vector dot-product is a function from two vectors to a scalar,
which generalizes to named tensors to give the ubiquitous contraction
operator. You can think of it as elementwise multiplication, then
summation over one axis: \[\begin{aligned} %A \ndot{height} B &= \
sum_i A_{\nidx{height}{i}} B_{\nidx{height}{i}} = \nmatrix{}{width}{
% 3\cdot2 + 1\cdot7 + 2\cdot1 & 1\cdot2 + 5\cdot7 + 6\cdot1 & 4\cdot2
+ 9\cdot7 + 5\cdot 1 %} \\ A \mathbin{\underset{\ensuremath{\mathsf
{width}}}{\odot}} C &= \sum_j A_{\ensuremath{\ensuremath{\mathsf
{width}}(j)}} B_{\ensuremath{\ensuremath{\mathsf{width}}(j)}} = \
ensuremath{\mathsf{height}}\begin{array}[b]{@{}c@{}}\ensuremath{\
mathsf{}}\\\begin{bmatrix} 3\cdot 1 + 1\cdot 4 + 4\cdot 1 \\ 1\cdot 1
+ 5\cdot 4 + 9\cdot 1 \\ 2\cdot 1 + 6\cdot 4 + 5\cdot 1 \end{bmatrix}
\end{array}.\end{aligned}\] See SS3.1.2 for more examples of what the
contraction operator can do.

Again, we can write multiple names to contract multiple axes at once.
An operator \(\odot\) with no axis name under it contracts zero axes
and is equivalent to elementwise multiplication, so we use \(\odot\)
for elementwise multiplication as well.

2.2.3 Renaming and reshaping

It's often useful to rename an axis (analogous to a transpose
operation in standard notation): \[A_{{\ensuremath{\mathsf{height}}\
rightarrow\ensuremath{\mathsf{height'}}}} = \ensuremath{\mathsf
{height'}}\begin{array}[b]{@{}c@{}}\ensuremath{\mathsf{width}}\\\
begin{bmatrix} 3 & 1 & 4 \\ 1 & 5 & 9 \\ 2 & 6 & 5 \\ \end{bmatrix}\
end{array}.\] We can also reshape two or more axes into one axis: \
[A_{{\ensuremath{\mathsf{(height,width)}}\rightarrow\ensuremath{\
mathsf{layer}}}} = \ensuremath{\mathsf{}}\begin{array}[b]{@{}c@{}}\
ensuremath{\mathsf{layer}}\\\begin{bmatrix} 3 & 1 & 4 & 1 & 5 & 9 & 2
& 6 & 5 \end{bmatrix}\end{array}\] The order of elements in the new
axis is undefined. If you need a particular order, you can write a
more specific definition.

3 Examples

In this section we give a series of examples illustrating how to use
named tensors in various situations, mostly related to machine
learning.

3.1 More tensor operations

3.1.1 Reductions

\[\begin{aligned} \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm
{min}}} A &= \min \{A_{\ensuremath{\ensuremath{\mathsf{ax}}(i)}} \mid
1 \leq i \leq n\} \\ \mathop{\underset{\ensuremath{\mathsf{ax}}}{\
mathrm{max}}} A &= \max \{A_{\ensuremath{\ensuremath{\mathsf{ax}}
(i)}} \mid 1 \leq i \leq n\} \\ \mathop{\underset{\ensuremath{\mathsf
{ax}}}{\mathrm{norm}}} A &= \sqrt{\sum\limits_{\ensuremath{\mathsf
{ax}}} A^2} \\ \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm
{mean}}} A &= \frac{1}{n} \sum\limits_{\ensuremath{\mathsf{ax}}} A \\
\mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{var}}} A &= \frac
{1}{n} \sum\limits_{\ensuremath{\mathsf{ax}}} (A - \mathop{\underset
{\ensuremath{\mathsf{ax}}}{\mathrm{mean}}} A)^2.\end{aligned}\] The \
(\min\) and \(\max\) operators are overloaded, as is the summation
operator defined above (SS3.1.1). If the operator is applied to a
tensor and has an axis under it, then it's a reduction performed over
the axis. But if it is applied to a set of tensors and has no axis
under it, then it's an elementwise operation performed over the set.

3.1.2 Contractions

The contraction operator can be used for many multiplication-like
operations.

\[\begin{aligned} u, v &\in \mathbb{R}^{\ensuremath{\mathsf{ax1}}} \\
x, y &\in \mathbb{R}^{\ensuremath{\mathsf{ax2}}} \\ A &\in \mathbb{R}
^{\ensuremath{\mathsf{ax1}} \times \ensuremath{\mathsf{ax2}}} \\ B &\
in \mathbb{R}^{\ensuremath{\mathsf{ax2}} \times \ensuremath{\mathsf
{ax3}}}\end{aligned}\]

\[\begin{aligned} u \mathbin{\underset{\ensuremath{\mathsf{ax1}}}{\
odot}} v &= \sum_i u_{\ensuremath{\ensuremath{\mathsf{ax1}}(i)}} v_{\
ensuremath{\ensuremath{\mathsf{ax1}}(i)}} && \text{inner product} \\
u \odot x &= \sum_{i,j} u_{\ensuremath{\ensuremath{\mathsf{ax1}}(i)}}
x_{\ensuremath{\ensuremath{\mathsf{ax2}}(j)}} && \text{outer product}
\\ A \mathbin{\underset{\ensuremath{\mathsf{ax2}}}{\odot}} x &= \
sum_j A_{\ensuremath{\ensuremath{\mathsf{ax2}}(j)}} x_{\ensuremath{\
ensuremath{\mathsf{ax2}}(j)}} && \text{matrix-vector multiplication}
\\ u \mathbin{\underset{\ensuremath{\mathsf{ax1}}}{\odot}} A &= \
sum_i u_{\ensuremath{\ensuremath{\mathsf{ax1}}(j)}} A_{\ensuremath{\
ensuremath{\mathsf{ax1}}(j)}} && \text{vector-matrix multiplication}
\\ A \mathbin{\underset{\ensuremath{\mathsf{ax2}}}{\odot}} B &= \
sum_j A_{\ensuremath{\ensuremath{\mathsf{ax2}}(j)}} \odot B_{\
ensuremath{\ensuremath{\mathsf{ax2}}(j)}} && \text{matrix-matrix
multiplication}\end{aligned}\]

3.1.3 Softmax and argmax

Most activation functions are elementwise operations (sigmoid, tanh,
ReLU), so they are straightforward to use in our notation; the
softmax, however, is interesting because it's defined as a function
from vectors to vectors: \[\begin{aligned} \mathop{\underset{\
ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} A &= \frac{\exp A}{\sum\
limits_{\ensuremath{\mathsf{ax}}} \exp A}.\end{aligned}\] As with
reductions, we write an axis below the softmax operator, but this
axis is retained in the output.

Closely related are argmax and argmin, which we define to compute
one-hot vectors with a one at the position containing the maximum or
minimum value. \[\begin{aligned} \mathop{\underset{\ensuremath{\
mathsf{ax}}}{\mathrm{argmax}}} A &= \lim_{\alpha \rightarrow \infty}
\mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} \alpha
A \\ \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{argmin}}} A
&= \lim_{\alpha \rightarrow -\infty} \mathop{\underset{\ensuremath{\
mathsf{ax}}}{\mathrm{softmax}}} \alpha A.\end{aligned}\]

3.2 Building blocks

3.2.1 Fully-connected layers

A feedforward neural network looks like this: \[\begin{aligned} X^0 &
\in \mathbb{R}^{\ensuremath{\mathsf{input}}} \\ X^1 &= \sigma(W^1 \
mathbin{\underset{\ensuremath{\mathsf{input}}}{\odot}} X^0 + b^1) & W
^1 &\in \mathbb{R}^{\ensuremath{\mathsf{hidden1}} \times \ensuremath
{\mathsf{input}}} & b^1 &\in \mathbb{R}^{\ensuremath{\mathsf
{hidden1}}} \\ X^2 &= \sigma(W^2 \mathbin{\underset{\ensuremath{\
mathsf{hidden1}}}{\odot}} X^1 + b^2) & W^2 &\in \mathbb{R}^{\
ensuremath{\mathsf{hidden2}} \times \ensuremath{\mathsf{hidden1}}} &
b^2 &\in \mathbb{R}^{\ensuremath{\mathsf{hidden2}}} \\ X^3 &= \sigma
(W^3 \mathbin{\underset{\ensuremath{\mathsf{hidden2}}}{\odot}} X^2 +
b^3) & W^3 &\in \mathbb{R}^{\ensuremath{\mathsf{output}} \times \
ensuremath{\mathsf{hidden2}}} & b^3 &\in \mathbb{R}^{\ensuremath{\
mathsf{output}}}\end{aligned}\] The layer sizes can be set by writing
\(|\ensuremath{\mathsf{input}}| = 100\), etc.

If you don't like repeating the equations for fully-connected layers,
you can put them inside a function: \[\begin{aligned} \text{FullConn}
^l(x) &= \sigma\left(W^l \mathbin{\underset{\ensuremath{\mathsf
{layer}}}{\odot}} x + b^l\right)_{{\ensuremath{\mathsf{layer'}}\
rightarrow\ensuremath{\mathsf{layer}}}}\end{aligned}\] where \[\begin
{aligned} W^l &\in \mathbb{R}^{\ensuremath{\ensuremath{\mathsf
{layer'}}[n_{l}]} \times \ensuremath{\ensuremath{\mathsf{layer}}[n_
{l-1}]}} \\ b^l &\in \mathbb{R}^{\ensuremath{\ensuremath{\mathsf
{layer'}}[n_l]}}.\end{aligned}\] A couple of things are new here.
First, \(\text{FullConn}^l\) encapsulates both the equation for layer
\(l\) as well as its parameters (analogous to what TensorFlow and
PyTorch call modules). Second, we chose to use the same axis name \(\
mathsf{layer}\) for all the layers (with different sizes \(n_l\)). So
\(\text{FullConn}^l\) temporarily computes its output over axis \(\
mathsf{layer'}\), then renames it back to \(\mathsf{layer}\).

Then the network can be defined like this: \[\begin{aligned} X^0 &\in
\mathbb{R}^{\ensuremath{\ensuremath{\mathsf{layer}}[n_0]}} \\ X^1 &=
\text{FullConn}^1(X^0) \\ X^2 &= \text{FullConn}^2(X^1) \\ X^3 &= \
text{FullConn}^3(X^2).\end{aligned}\]

3.2.2 Recurrent neural networks

As a second example, let's define a simple (Elman) RNN. This is
similar to the feedforward network, except that the number of
timesteps is variable and they all share parameters. \[\begin
{aligned} x^{t} &\in \mathbb{R}^{\ensuremath{\mathsf{input}}} & t &=
1, \ldots, n \\ W^{\text{h}} &\in \mathbb{R}^{\ensuremath{\mathsf
{hidden}} \times \ensuremath{\mathsf{hidden'}}} & |\ensuremath{\
mathsf{hidden}}| &= |\ensuremath{\mathsf{hidden'}}| \\ W^{\text{i}} &
\in \mathbb{R}^{\ensuremath{\mathsf{input}} \times \ensuremath{\
mathsf{hidden'}}} \\ b &\in \mathbb{R}^{\ensuremath{\mathsf
{hidden'}}} \\ h^{0} &\in \mathbb{R}^{\ensuremath{\mathsf{hidden}}} \
\ h^{t} &= \sigma\left( W^{\text{h}} \mathbin{\underset{\ensuremath{\
mathsf{hidden}}}{\odot}} h^{t-1} + W^{\text{i}} \mathbin{\underset{\
ensuremath{\mathsf{input}}}{\odot}} x^{t} + b \right)_{{\ensuremath{\
mathsf{hidden'}}\rightarrow\ensuremath{\mathsf{hidden}}}} & t &= 1, \
ldots, n\end{aligned}\]

3.2.3 Attention

In the introduction (SS1), we mentioned some difficulties in
interpreting the equation for attention as it's usually written. In
our notation, it looks like this: \[\begin{aligned} \text{Attention}
\colon \mathbb{R}^{\ensuremath{\mathsf{key}}} \times \mathbb{R}^{\
ensuremath{\mathsf{seq}} \times\ensuremath{\mathsf{key}}} \times \
mathbb{R}^{\ensuremath{\mathsf{seq}} \times\ensuremath{\mathsf{val}}}
&\rightarrow \mathbb{R}^{\ensuremath{\mathsf{val}}} \\ \text
{Attention}(Q,K,V) &= \mathop{\underset{\ensuremath{\mathsf{seq}}}{\
mathrm{softmax}}} \left( \frac{Q \mathbin{\underset{\ensuremath{\
mathsf{key}}}{\odot}} K}{\sqrt{|\ensuremath{\mathsf{key}}|}} \right)
\mathbin{\underset{\ensuremath{\mathsf{seq}}}{\odot}} V.\end{aligned}
\]

This equation is slightly different from the one in the introduction.
The previous definition computed an output sequence over axis \(\
mathsf{seq'}\), but this definition computes a single value. If we
want a sequence, we can just give \(Q\) a \(\mathsf{seq'}\) axis (or
some other name), and the function will compute an output sequence.
Furthermore, if we give \(Q\), \(K\), and \(V\) a \(\mathsf{heads}\)
axis for multiple attention heads, then the function will compute
multi-head attention.

Sometimes we need to apply a mask to keep from attending to certain
positions. \[\begin{aligned} \text{Attention} \colon \mathbb{R}^{\
ensuremath{\mathsf{key}}} \times \mathbb{R}^{\ensuremath{\mathsf
{seq}} \times\ensuremath{\mathsf{key}}} \times \mathbb{R}^{\
ensuremath{\mathsf{seq}} \times\ensuremath{\mathsf{val}}} \times \
mathbb{R}^{\ensuremath{\mathsf{seq}}} &\rightarrow \mathbb{R}^{\
ensuremath{\mathsf{val}}} \\ \text{Attention}(Q, K, V, M) &= \mathop
{\underset{\ensuremath{\mathsf{seq}}}{\mathrm{softmax}}} \left( \frac
{Q \mathbin{\underset{\ensuremath{\mathsf{key}}}{\odot}} K}{\sqrt{|\
ensuremath{\mathsf{key}}|}} + M \right) \mathbin{\underset{\
ensuremath{\mathsf{seq}}}{\odot}} V.\end{aligned}\]

3.2.4 Convolution

A 1-dimensional convolution can be written using contractions: \[\
begin{aligned} \text{Conv1d} \colon \mathbb{R}^{\ensuremath{\mathsf
{chans}} \times \ensuremath{\ensuremath{\mathsf{seq}}[n]}} &\
rightarrow \mathbb{R}^{\ensuremath{\ensuremath{\mathsf{seq}}[n']}} \\
\text{Conv1d}(X; W, b) &= [W \mathbin{\underset{\ensuremath{\mathsf
{chans,kernel}}}{\odot}} C \mathbin{\underset{\ensuremath{\mathsf
{seq}}}{\odot}} X + b]_{{\ensuremath{\mathsf{out}}\rightarrow\
ensuremath{\mathsf{seq}}}}\end{aligned}\] where \[\begin{aligned} W &
\in \mathbb{R}^{\ensuremath{\mathsf{chans}} \times \ensuremath{\
mathsf{kernel}}} \\ b &\in \mathbb{R}\\ n' &= n-|\ensuremath{\mathsf
{kernel}}|+1 \\ C &\in \mathbb{R}^{\ensuremath{\ensuremath{\mathsf
{out}}[n']} \times \ensuremath{\mathsf{kernel}} \times \ensuremath{\
ensuremath{\mathsf{seq}}[n]}} \\ C_{\ensuremath{\ensuremath{\mathsf
{out}}(o)},\ensuremath{\ensuremath{\mathsf{kernel}}(k)},\ensuremath{\
ensuremath{\mathsf{seq}}(i)}} &= \delta(o+k-1,i).\end{aligned}\] This
computes a single output channel, but we can get multiple output
channels by giving \(W\) and \(b\) a \(\mathsf{chans'}\) axis (or
some other name).

We can use the same \(C\) to define a 2-dimensional convolution: \[\
begin{aligned} \text{Conv2d} \colon \mathbb{R}^{\ensuremath{\mathsf
{chans}} \times \ensuremath{\ensuremath{\mathsf{height}}[h]} \times \
ensuremath{\ensuremath{\mathsf{width}}[w]}} &\rightarrow \mathbb{R}^
{\ensuremath{\ensuremath{\mathsf{height}}[h']} \times \ensuremath{\
ensuremath{\mathsf{width}}[w']}} \\ \text{Conv2d}(X; W, b) &= [W \
mathbin{\underset{\ensuremath{\mathsf{chans, kh, kw}}}{\odot}} C^2 \
mathbin{\underset{\ensuremath{\mathsf{height,width}}}{\odot}} X + b]_
{\substack{{\ensuremath{\mathsf{oh}}\rightarrow\ensuremath{\mathsf
{height}}}\\{\ensuremath{\mathsf{ow}}\rightarrow\ensuremath{\mathsf
{width}}}}}\end{aligned}\] where \[\begin{aligned} W &\in \mathbb{R}^
{\ensuremath{\mathsf{chans}} \times \ensuremath{\mathsf{kh}} \times \
ensuremath{\mathsf{kw}}} \\ b &\in \mathbb{R}\\ h' &= h-|\ensuremath
{\mathsf{kh}}|+1 \\ w' &= w-|\ensuremath{\mathsf{kw}}|+1 \\ C^2 &= C_
{\substack{{\ensuremath{\mathsf{out}}\rightarrow\ensuremath{\mathsf
{oh}}} \\ {\ensuremath{\mathsf{kernel}}\rightarrow\ensuremath{\mathsf
{kh}}} \\ {\ensuremath{\mathsf{seq}}\rightarrow\ensuremath{\mathsf
{height}}}}} \odot C_{\substack{{\ensuremath{\mathsf{out}}\rightarrow
\ensuremath{\mathsf{ow}}} \\ {\ensuremath{\mathsf{kernel}}\rightarrow
\ensuremath{\mathsf{kw}}} \\ {\ensuremath{\mathsf{seq}}\rightarrow\
ensuremath{\mathsf{width}}}}}.\end{aligned}\]

3.2.5 Max pooling

\[\begin{aligned} \text{MaxPool1d}_{k} \colon \mathbb{R}^{\ensuremath
{\ensuremath{\mathsf{seq}}[n]}} &\rightarrow \mathbb{R}^{\ensuremath
{\ensuremath{\mathsf{seq}}[n/k]}} \\ \text{MaxPool1d}_{k}(X) &= \
mathop{\underset{\ensuremath{\mathsf{k}}}{\mathrm{max}}} U\end
{aligned}\] where \[\begin{aligned} U &\in \mathbb{R}^{\ensuremath{\
ensuremath{\mathsf{seq}}[n / k]} \times \ensuremath{\ensuremath{\
mathsf{k}}[k]}} \\ U_{\ensuremath{\ensuremath{\mathsf{seq}}(i)}, \
ensuremath{\ensuremath{\mathsf{k}}(di)}} & = X_{\ensuremath{\
ensuremath{\mathsf{seq}}(i \times k + di -1)}}.\end{aligned}\]

\[\begin{aligned} \text{MaxPool2d}_{kh,kw} \colon \mathbb{R}^{\
ensuremath{\ensuremath{\mathsf{height}}[h]} \times \ensuremath{\
ensuremath{\mathsf{width}}[w]}} &\rightarrow \mathbb{R}^{\ensuremath
{\ensuremath{\mathsf{height}}[h/kh]} \times \ensuremath{\ensuremath{\
mathsf{width}}[w/kw]}} \\ \text{MaxPool2d}_{kh,hw}(X) &= \mathop{\
underset{\ensuremath{\mathsf{kh, kw}}}{\mathrm{max}}} U\end{aligned}
\] where \[\begin{aligned} U &\in \mathbb{R}^{\ensuremath{\ensuremath
{\mathsf{height}}[h / kh]} \times \ensuremath{\ensuremath{\mathsf
{width}}[w / kw]} \times \ensuremath{\ensuremath{\mathsf{kh}}[kh]} \
times \ensuremath{\ensuremath{\mathsf{kw}}[kw]}} \\ U_{\ensuremath{\
ensuremath{\mathsf{height}}(i)}, \ensuremath{\ensuremath{\mathsf
{width}}(j)}, \ensuremath{\ensuremath{\mathsf{kh}}(di)}, \ensuremath
{\ensuremath{\mathsf{kw}}(dj)}} & = X_{\ensuremath{\ensuremath{\
mathsf{height}}(i \times kh + di -1)}, \ensuremath{\ensuremath{\
mathsf{width}}(j \times kw + dj -1)}}.\end{aligned}\]

3.2.6 Normalization layers

Batch, instance, and layer normalization are often informally
described using the same equation, but they each correspond to very
different functions. They differ by which axes are normalized.

We can define a single generic normalization layer: \[\begin{aligned}
\mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{XNorm}}} \colon \
mathbb{R}^{\ensuremath{\mathsf{ax}}} &\rightarrow \mathbb{R}^{\
ensuremath{\mathsf{ax}}} \\ \mathop{\underset{\ensuremath{\mathsf
{ax}}}{\mathrm{XNorm}}}(X; \gamma, \beta, \epsilon) &= \frac{X - \
mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{mean}}}(X)}{\sqrt
{\mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{var}}}(X)} + \
epsilon} \odot \gamma + \beta\end{aligned}\] where \[\begin{aligned}
\gamma, \beta &\in \mathbb{R}^{\ensuremath{\mathsf{ax}}} \\ \epsilon
&> 0.\end{aligned}\]

Now, suppose that the input has three axes: \[\begin{aligned} X &\in
\mathbb{R}^{{\ensuremath{\mathsf{batch}} \times \ensuremath{\mathsf
{chans}} \times \ensuremath{\mathsf{layer}}}}\end{aligned}\] Then the
three kinds of normalization layers can be written as: \[\begin
{aligned} Y &= \mathop{\underset{\ensuremath{\mathsf{batch}}}{\mathrm
{XNorm}}}(X; \gamma, \beta) && \text{batch normalization} \\ Y &= \
mathop{\underset{\ensuremath{\mathsf{layer}}}{\mathrm{XNorm}}}(X; \
gamma, \beta) && \text{instance normalization} \\ Y &= \mathop{\
underset{\ensuremath{\mathsf{layer,chans}}}{\mathrm{XNorm}}}(X; \
gamma, \beta) && \text{layer normalization}\end{aligned}\]

3.3 Transformer

We define a Transformer used autoregressively as a language model.
The input is a sequence of one-hot vectors, from which we compute
word embeddings and positional encodings: \[\begin{aligned} I &\in \
{0, 1\}^{\ensuremath{\mathsf{seq}} \times \ensuremath{\mathsf
{vocab}}} & \sum\limits_{\ensuremath{\mathsf{vocab}}} I &= 1 \\ W &=
(E \mathbin{\underset{\ensuremath{\mathsf{vocab}}}{\odot}} I)\sqrt{|\
ensuremath{\mathsf{layer}}|} & E &\in \mathbb{R}^{\ensuremath{\mathsf
{vocab}} \times \ensuremath{\mathsf{layer}}} \\ P &\in \mathbb{R}^{\
ensuremath{\mathsf{seq}} \times \ensuremath{\mathsf{layer}}} \\ P_{\
ensuremath{\ensuremath{\mathsf{seq}}(p)}, \ensuremath{\ensuremath{\
mathsf{layer}}(i)}} &= \begin{cases} \sin((p-1) / 10000^{(i-1) / |\
ensuremath{\mathsf{layer}}|}) & \text{$i$ odd} \\ \cos((p-1) / 10000^
{(i-2) / |\ensuremath{\mathsf{layer}}|}) & \text{$i$ even.} \end
{cases}\end{aligned}\]

Then we use \(L\) layers of self-attention and feed-forward neural
networks: \[\begin{aligned} X^0 &= W+P \\ T^1 &= \text{LayerNorm}^1(\
text{SelfAtt}^1(X^0)) + X^0\\ X^1 &= \text{LayerNorm}^{1'}(\text{FFN}
^1(T^1)) + T^1\\ &\vdotswithin{=} \\ T^{L} &= \text{LayerNorm}^L(\
text{SelfAtt}^L(X^{L-1})) + X^{L-1}\\ X^{L} &= \text{LayerNorm}^{L'}
(\text{FFN}^L(T^L)) + T^L\\ O &= \mathop{\underset{\ensuremath{\
mathsf{vocab}}}{\mathrm{softmax}}}(E \mathbin{\underset{\ensuremath{\
mathsf{layer}}}{\odot}} X^L)\end{aligned}\] where \(\text{LayerNorm}
\), \(\text{SelfAtt}\) and \(\text{FFN}\) are defined below.

Layer normalization (\(l = 1, 1', \ldots, L, L'\)): \[\begin{aligned}
\text{LayerNorm}^l \colon \mathbb{R}^{\ensuremath{\mathsf{layer}}} &\
rightarrow \mathbb{R}^{\ensuremath{\mathsf{layer}}} \\ \text
{LayerNorm}^l(X) &= \mathop{\underset{\ensuremath{\mathsf{layer}}}{\
mathrm{XNorm}}}(X; \beta^l, \gamma^l).\end{aligned}\]

We defined attention in SS3.2.3; the Transformer uses multi-head
self-attention, in which queries, keys, and values are all computed
from the same sequence. \[\begin{aligned} \text{SelfAtt}^l \colon \
mathbb{R}^{\ensuremath{\mathsf{seq}} \times \ensuremath{\mathsf
{layer}}} &\rightarrow \mathbb{R}^{\ensuremath{\mathsf{seq}} \times \
ensuremath{\mathsf{layer}}} \\ \text{SelfAtt}^l(X) &= Y\end{aligned}
\] where \[\begin{aligned} |\ensuremath{\mathsf{seq}}| &= |\
ensuremath{\mathsf{seq'}}| \\ |\ensuremath{\mathsf{key}}| = |\
ensuremath{\mathsf{val}}| &= |\ensuremath{\mathsf{layer}}|/|\
ensuremath{\mathsf{heads}}| \\ Q &= W^{l,Q} \mathbin{\underset{\
ensuremath{\mathsf{layer}}}{\odot}} X_{{\ensuremath{\mathsf{seq}}\
rightarrow\ensuremath{\mathsf{seq'}}}} & W^{l,Q} &\in \mathbb{R}^{\
ensuremath{\mathsf{heads}} \times \ensuremath{\mathsf{layer}} \times
\ensuremath{\mathsf{key}}} \\ K &= W^{l,K} \mathbin{\underset{\
ensuremath{\mathsf{layer}}}{\odot}} X & W^{l,K} &\in \mathbb{R}^{\
ensuremath{\mathsf{heads}} \times \ensuremath{\mathsf{layer}} \times
\ensuremath{\mathsf{key}}} \\ V &= W^{l,V} \mathbin{\underset{\
ensuremath{\mathsf{layer}}}{\odot}} X & W^{l,V} &\in \mathbb{R}^{\
ensuremath{\mathsf{heads}} \times \ensuremath{\mathsf{layer}} \times
\ensuremath{\mathsf{val}}} \\ M & \in \mathbb{R}^{\ensuremath{\mathsf
{seq}} \times \ensuremath{\mathsf{seq'}}} \\ M_{\ensuremath{\
ensuremath{\mathsf{seq}}(i)}, \ensuremath{\ensuremath{\mathsf{seq'}}
(j)}} &= \begin{cases} 0 & i \leq j\\ -\infty & \text{otherwise} \end
{cases} \\ Y &= W^{l,O} \mathbin{\underset{\ensuremath{\mathsf
{heads,val}}}{\odot}} \text{Attention}(Q, K, V, M)_{{\ensuremath{\
mathsf{seq'}}\rightarrow\ensuremath{\mathsf{seq}}}} & W^{l,O} &\in \
mathbb{R}^{\ensuremath{\mathsf{heads}} \times \ensuremath{\mathsf
{val}} \times \ensuremath{\mathsf{layer}}}\end{aligned}\]

Feedforward neural networks: \[\begin{aligned} \text{FFN}^l \colon \
mathbb{R}^{\ensuremath{\mathsf{layer}}} &\rightarrow \mathbb{R}^{\
ensuremath{\mathsf{layer}}} \\ \text{FFN}^l(X) &= X^2\end{aligned}\]
where \[\begin{aligned} X^1 &= \text{relu}(W^{l,1} \mathbin{\underset
{\ensuremath{\mathsf{layer}}}{\odot}} X + b^{l,1}) & W^{l,1} &\in \
mathbb{R}^{\ensuremath{\mathsf{hidden}} \times \ensuremath{\mathsf
{layer}}} & b^{l,1} &\in \mathbb{R}^{\ensuremath{\mathsf{hidden}}} \\
X^2 &= \text{relu}(W^{l,2} \mathbin{\underset{\ensuremath{\mathsf
{hidden}}}{\odot}} X^1 + b^{l,2}) & W^{l,2} &\in \mathbb{R}^{\
ensuremath{\mathsf{layer}} \times \ensuremath{\mathsf{hidden}}} & b^
{l,2} &\in \mathbb{R}^{\ensuremath{\mathsf{hidden}}}.\end{aligned}\]

3.4 LeNet

\[\begin{aligned} X^0 &\in \mathbb{R}^{\ensuremath{\mathsf{batch}} \
times \ensuremath{\ensuremath{\mathsf{chans}}[c_0]} \times \
ensuremath{\mathsf{height}} \times \ensuremath{\mathsf{width}}} \\ T^
1 &= \text{relu}(\text{Conv}^1(X^0)) \\ X^1 &= \text{MaxPool}^1(T^1)
\\ T^2 &= \text{relu}(\text{Conv}^2(X^1)) \\ X^2 &= \text{MaxPool}^2
(T^2)_{{\ensuremath{\mathsf{(height,width,chans)}}\rightarrow\
ensuremath{\mathsf{layer}}}} \\ X^3 &= \text{relu}(W^3 \mathbin{\
underset{\ensuremath{\mathsf{layer}}}{\odot}} X^2 + b^3) & W^3 &\in \
mathbb{R}^{\ensuremath{\mathsf{hidden}} \times \ensuremath{\mathsf
{layer}}} & b^3 &\in \mathbb{R}^{\ensuremath{\mathsf{hidden}}} \\ O &
= \mathop{\underset{\ensuremath{\mathsf{classes}}}{\mathrm{softmax}}}
(W^4 \mathbin{\underset{\ensuremath{\mathsf{hidden}}}{\odot}} X^3 + b
^4) & W^4 &\in \mathbb{R}^{\ensuremath{\mathsf{classes}} \times \
ensuremath{\mathsf{hidden}}} & b^4 &\in \mathbb{R}^{\ensuremath{\
mathsf{classes}}}\end{aligned}\] As an alternative to the flattening
operation in the equation for \(X^2\), we could have written \[\begin
{aligned} X^2 &= \text{MaxPool}^2(T^2) \\ X^3 &= \text{relu}(W^3 \
mathbin{\underset{\ensuremath{\mathsf{height,width,chans}}}{\odot}} X
^2 + b^3) & W^3 &\in \mathbb{R}^{\ensuremath{\mathsf{hidden}} \times
\ensuremath{\mathsf{height}} \times \ensuremath{\mathsf{width}} \
times \ensuremath{\mathsf{chans}}}.\end{aligned}\]

The convolution and pooling operations are defined as follows: \[\
begin{aligned} \text{Conv}^l(X) &= \text{Conv2d}(X; W^l, b^l)_{{\
ensuremath{\mathsf{chans'}}\rightarrow\ensuremath{\mathsf{chans}}}}\
end{aligned}\] where \[\begin{aligned} W^l & \in \mathbb{R}^{\
ensuremath{\ensuremath{\mathsf{chans'}}[c_{l}]} \times \ensuremath{\
ensuremath{\mathsf{chans}}[c_{l-1}]} \times \ensuremath{\ensuremath{\
mathsf{kh}}[kh_l]} \times \ensuremath{\ensuremath{\mathsf{kw}}
[kw_l]}} \\ b^l &\in \mathbb{R}^{\ensuremath{\ensuremath{\mathsf
{chans'}}[c_{l}]}}\end{aligned}\] and \[\begin{aligned} \text
{MaxPool}^l(X) &= \text{MaxPol2d}_{ph^l,ph^l}(X).\end{aligned}\]

3.5 Other examples

3.5.1 Discrete random variables

Named axes are very helpful for working with discrete random
variables, because each random variable can be represented by an axis
with the same name. For instance, if \(\ensuremath{\mathsf{A}}\) and 
\(\ensuremath{\mathsf{B}}\) are random variables, we can treat \(p(\
ensuremath{\mathsf{B}} \mid \ensuremath{\mathsf{A}})\) and \(p(\
ensuremath{\mathsf{A}})\) as tensors: \[\begin{aligned} p(\ensuremath
{\mathsf{B}} \mid \ensuremath{\mathsf{A}}) &\in [0, 1]^{\ensuremath{\
mathsf{A}} \times \ensuremath{\mathsf{B}}} & \sum\limits_{\ensuremath
{\mathsf{B}}} p(\ensuremath{\mathsf{B}}\mid \ensuremath{\mathsf{A}})
&= 1 \\ p(\ensuremath{\mathsf{A}}) &\in [0, 1]^{\ensuremath{\mathsf
{A}}} & \sum\limits_{\ensuremath{\mathsf{A}}} p(\ensuremath{\mathsf
{A}}) &= 1\end{aligned}\] Then many common operations on probability
distributions can be expressed in terms of tensor operations: \[\
begin{aligned} p(\ensuremath{\mathsf{A}}, \ensuremath{\mathsf{B}}) &=
p(\ensuremath{\mathsf{B}} \mid \ensuremath{\mathsf{A}}) \odot p(\
ensuremath{\mathsf{A}}) && \text{chain rule}\\ p(\ensuremath{\mathsf
{B}}) &= \sum\limits_{\ensuremath{\mathsf{A}}} p(\ensuremath{\mathsf
{A}}, \ensuremath{\mathsf{B}}) = p(\ensuremath{\mathsf{B}} \mid \
ensuremath{\mathsf{A}}) \mathbin{\underset{\ensuremath{\mathsf{A}}}{\
odot}} p(\ensuremath{\mathsf{A}}) && \text{marginalization} \\ p(\
ensuremath{\mathsf{A}} \mid \ensuremath{\mathsf{B}}) &= \frac{p(\
ensuremath{\mathsf{A}}, \ensuremath{\mathsf{B}})}{p(\ensuremath{\
mathsf{B}})} = \frac{p(\ensuremath{\mathsf{B}} \mid \ensuremath{\
mathsf{A}}) \odot p(\ensuremath{\mathsf{A}})}{p(\ensuremath{\mathsf
{B}} \mid \ensuremath{\mathsf{A}}) \mathbin{\underset{\ensuremath{\
mathsf{A}}}{\odot}} p(\ensuremath{\mathsf{A}})}. && \text{Bayes'
rule}\end{aligned}\]

3.5.2 Continuous bag of words

A continuous bag-of-words model classifies by summing up the
embeddings of a sequence of words \(X\) and then projecting them to
the space of classes.

\[\begin{aligned} \text{CBOW} \colon \{0, 1\}^{\ensuremath{\mathsf
{seq}} \times \ensuremath{\mathsf{vocab}}} &\rightarrow \mathbb{R}^{\
ensuremath{\mathsf{seq}} \times \ensuremath{\mathsf{classes}}} \\ \
text{CBOW}(X; E, W) &= \mathop{\underset{\ensuremath{\mathsf{class}}}
{\mathrm{softmax}}} (W \mathbin{\underset{\ensuremath{\mathsf
{hidden}}}{\odot}} E \mathbin{\underset{\ensuremath{\mathsf{vocab}}}
{\odot}} X)\end{aligned}\] where \[\begin{aligned} \sum\limits_{\
ensuremath{\mathsf{vocab}}} X &= 1 \\ E &\in \mathbb{R}^{\ensuremath
{\mathsf{vocab}} \times \ensuremath{\mathsf{hidden}}} \\ W &\in \
mathbb{R}^{\ensuremath{\mathsf{classes}} \times \ensuremath{\mathsf
{hidden}}}.\end{aligned}\] Here, the two contractions can be done in
either order, so we leave the parentheses off.

3.5.3 Sudoku ILP

Sudoku puzzles can be represented as binary tiled tensors. Given a
grid we can check that it is valid by converting it to a grid of
grids. Constraints then ensure that there is one digit per row, per
column and per sub-box.

\[\begin{aligned} \text{check} \colon \{0, 1\}^{\ensuremath{\
ensuremath{\mathsf{height}}[9]} \times \ensuremath{\ensuremath{\
mathsf{width}}[9]} \times \ensuremath{\ensuremath{\mathsf{assign}}
[9]}} &\rightarrow \{0, 1\} \\ \text{check}(X) &= \mathbb{I}\left[\
begin{aligned} \sum\limits_{\ensuremath{\mathsf{assign}}} X = 1 &\
land \sum\limits_{\ensuremath{\mathsf{height, width}}} Y = 1 \land {}
\\ \sum\limits_{\ensuremath{\mathsf{height}}} X = 1 &\land \sum\
limits_{\ensuremath{\mathsf{width}}} X = 1 \end{aligned}\right]\end
{aligned}\] where \[\begin{aligned} Y &\in \{0, 1\}^{\ensuremath{\
ensuremath{\mathsf{height'}}[3]} \times \ensuremath{\ensuremath{\
mathsf{width'}}[3]} \times \ensuremath{\ensuremath{\mathsf{height}}
[3]} \times \ensuremath{\ensuremath{\mathsf{width}}[3]} \times \
ensuremath{\ensuremath{\mathsf{assign}}[9]}} \\ Y_{\ensuremath{\
ensuremath{\mathsf{height'}}(h')}, \ensuremath{\ensuremath{\mathsf
{height}}(h)}, \ensuremath{\ensuremath{\mathsf{width'}}(w')}, \
ensuremath{\ensuremath{\mathsf{width}}(w)}} &= X_{\ensuremath{\
ensuremath{\mathsf{height}}(3h' + h-1)}, \ensuremath{\ensuremath{\
mathsf{width}}(3 w' + w-1)}}.\end{aligned}\]

3.5.4 \(K\)-means clustering

The following equations define one step of \(k\)-means clustering.
Given a set of points \(X\) and an initial set of cluster centers \(C
\), \[\begin{aligned} X &\in \mathbb{R}^{\ensuremath{\mathsf{batch}}
\times \ensuremath{\mathsf{space}}} \\ C &\in \mathbb{R}^{\ensuremath
{\mathsf{clusters}} \times \ensuremath{\mathsf{space}}}\end{aligned}
\] we repeat the following update: Compute cluster assignments \[\
begin{aligned} Q &= \mathop{\underset{\ensuremath{\mathsf{clusters}}}
{\mathrm{argmin}}} \mathop{\underset{\ensuremath{\mathsf{space}}}{\
mathrm{norm}}}(C-X)\end{aligned}\] then recompute the cluster
centers: \[C \leftarrow \sum\limits_{\ensuremath{\mathsf{batch}}} \
frac{Q \odot X}{Q}.\]

3.5.5 Beam search

Beam search is a commonly used approach for approximate discrete
search. Here \(H\) is the score of each element in the beam, \(S\) is
the state of each element in the beam, and \(f\) is an update
function that returns the score of each state transition. \[\begin
{aligned} H &\in \mathbb{R}^{\ensuremath{\mathsf{beam}}} \\ S &\in \
{0, 1\}^{\ensuremath{\mathsf{beam}} \times \ensuremath{\mathsf
{state}}} & \sum\limits_{\ensuremath{\mathsf{state}}} S &= 1 \\ f &\
colon \{0, 1\}^{\ensuremath{\mathsf{state}}} \rightarrow \mathbb{R}^
{\ensuremath{\mathsf{state}}} \\\end{aligned}\] Then we repeat the
following update: \[\begin{aligned} H' &= \mathop{\underset{\
ensuremath{\mathsf{beam}}}{\mathrm{max}}} (H \odot f(S)) \\ H &\
leftarrow \mathop{\underset{\ensuremath{\mathsf{state,beam}}}{\mathrm
{maxk}}} H' \\ S &\leftarrow \mathop{\underset{\ensuremath{\mathsf
{state,beam}}}{\mathrm{argmaxk}}} H'\end{aligned}\] where \[\begin
{aligned} \mathop{\underset{\ensuremath{\mathsf{ax,k}}}{\mathrm
{maxk}}} \colon \mathbb{R}^{\ensuremath{\mathsf{ax}}} &\rightarrow \
mathbb{R}^{\ensuremath{\mathsf{k}}} \\ \mathop{\underset{\ensuremath
{\mathsf{ax,k}}}{\mathrm{argmaxk}}} \colon \mathbb{R}^{\ensuremath{\
mathsf{ax}}} &\rightarrow \{0,1\}^{\ensuremath{\mathsf{ax}},\
ensuremath{\mathsf{k}}}\end{aligned}\] are defined such that \([\
mathop{\underset{\ensuremath{\mathsf{ax,k}}}{\mathrm{maxk}}} A]_{\
ensuremath{\ensuremath{\mathsf{k}}(i)}}\) is the \(i\)-th largest
value along axis \(\ensuremath{\mathsf{ax}}\) and \(A \mathbin{\
underset{\ensuremath{\mathsf{ax}}}{\odot}} (\mathop{\underset{\
ensuremath{\mathsf{ax,k}}}{\mathrm{argmaxk}}}{A}) = \mathop{\underset
{\ensuremath{\mathsf{ax,k}}}{\mathrm{max}}} A\).

We can add a \(\mathsf{batch}\) axis to \(H\) and \(S\) and the above
equations will work unchanged.

3.5.6 Multivariate normal distribution

To define a multivariate normal distribution, we need some matrix
operations. These have two axis names written under them, for rows
and columns, respectively. Determinant and inverse have the following
signatures: \[\begin{aligned} \mathop{\underset{\ensuremath{\mathsf
{ax1,ax2}}}{\mathrm{det}}} \colon F^{\ensuremath{\ensuremath{\mathsf
{ax1}}[n]} \times \ensuremath{\ensuremath{\mathsf{ax2}}[n]}} &\
rightarrow F \\ \mathop{\underset{\ensuremath{\mathsf{ax1,ax2}}}{\
mathrm{inv}}} \colon F^{\ensuremath{\ensuremath{\mathsf{ax1}}[n]} \
times \ensuremath{\ensuremath{\mathsf{ax2}}[n]}} &\rightarrow F^{\
ensuremath{\ensuremath{\mathsf{ax1}}[n]} \times \ensuremath{\
ensuremath{\mathsf{ax2}}[n]}}.\end{aligned}\] (We write \(\text{inv}
\) instead of \(\cdot^{-1}\) because there's no way to write axis
names under the latter.)

In our notation, the application of a bilinear form is more verbose
than the standard notation (\((X-\mu)^\top \Sigma^{-1} (X-\mu)\)),
but also makes it look more like a function of two arguments (and
would generalize to three or more arguments).

\[\begin{aligned} \mathcal{N} \colon \mathbb{R}^{\ensuremath{\mathsf
{d}}} &\rightarrow \mathbb{R}\\ \mathcal{N}(X; \mu, \Sigma) &= \frac
{\exp\left(-\frac{1}{2} \left(\mathop{\underset{\ensuremath{\mathsf
{d1, d2}}}{\mathrm{inv}}} \Sigma\right) \mathbin{\underset{\
ensuremath{\mathsf{d1,d2}}}{\odot}} \left([X - \mu]_{{\ensuremath{\
mathsf{d}}\rightarrow\ensuremath{\mathsf{d1}}}} \odot [X - \mu]_{{\
ensuremath{\mathsf{d}}\rightarrow\ensuremath{\mathsf{d2}}}} \right) \
right)}{\sqrt{(2 \pi)^{|\ensuremath{\mathsf{d}}|} \mathop{\underset{\
ensuremath{\mathsf{d1, d2}}}{\mathrm{det}}} \Sigma}}\end{aligned}\]
where \[\begin{aligned} |\ensuremath{\mathsf{d}}| &= |\ensuremath{\
mathsf{d1}}| = |\ensuremath{\mathsf{d2}}| \\ \mu &\in \mathbb{R}^{\
ensuremath{\mathsf{d}}} \\ \Sigma & \in \mathbb{R}^{\ensuremath{\
mathsf{d1}} \times \ensuremath{\mathsf{d2}}}.\end{aligned}\]

4 LaTeX Macros

Many of the LaTeX macros used in this document are available in the
style file https://namedtensor.github.io/namedtensor.sty. To use it,
put

    \usepackage{namedtensor}

in the preamble of your LaTeX source file (after \documentclass
{article} but before \begin{document}).

The style file contains a small number of macros:

  * Basics

      + Use \name{foo} to write an axis name: \(\ensuremath{\mathsf
        {foo}}\).

      + Use \mathbb{R}^{\nset{foo}{2}} to write a set of tensors: \(\
        mathbb{R}^{\ensuremath{\ensuremath{\mathsf{foo}}[2]}}\).

      + Use A_{\nidx{foo}{1}} to index a tensor: \(A_{\ensuremath{\
        ensuremath{\mathsf{foo}}(1)}}\).

      + Use A_{\nmov{foo}{bar}} for renaming: \(A_{{\ensuremath{\
        mathsf{foo}}\rightarrow\ensuremath{\mathsf{bar}}}}\).

  * Binary operators

      + Use A \ndot{foo} B for contraction: \(A \mathbin{\underset{\
        ensuremath{\mathsf{foo}}}{\odot}} B\).

      + Use A \ncat{foo} B for concatenation: \(A \mathbin{\underset
        {\ensuremath{\mathsf{foo}}}{\oplus}} B\).

      + In general, you can use \nbin to make a new binary operator
        with a name under it: A \nbin{foo}{\star} B gives you \(A \
        mathbin{\underset{\ensuremath{\mathsf{foo}}}{\star}} B\).

  * Functions

      + Use \nsum{foo} A for summation: \(\sum\limits_{\ensuremath{\
        mathsf{foo}}} A\).

      + In general, you can use \nfun to make a function with a name
        under it: \nfun{foo}{qux} A gives you \(\mathop{\underset{\
        ensuremath{\mathsf{foo}}}{\mathrm{qux}}} A\).

5 Formal Definitions

5.1 Records and shapes

A named index is a pair, written \(\ensuremath{\ensuremath{\mathsf
{ax}}(i)}\), where \(\ensuremath{\mathsf{ax}}\) is a name and \(i\)
is usually a natural number. We write both names and variables
ranging over names using sans-serif font.

A record is a set of named indices \(\{\ensuremath{\ensuremath{\
mathsf{ax_\text{$1$}}}(i_1)}, \ldots, \ensuremath{\ensuremath{\mathsf
{ax_\text{$r$}}}(i_r)}\}\), where \(\ensuremath{\mathsf{ax_\text
{$1$}}}, \ldots \ensuremath{\mathsf{ax_\text{$r$}}}\) are pairwise
distinct names.

An axis is a pair, written \(\ensuremath{\ensuremath{\mathsf{ax}}[I]}
\), where \(\ensuremath{\mathsf{ax}}\) is a name and \(I\) is a set
of indices. We deal with axes of the form \(\ensuremath{\ensuremath{\
mathsf{ax}}[[n]]}\) (that is, \(\ensuremath{\ensuremath{\mathsf{ax}}
[\{1, \ldots, n\}]}\)) so frequently that we abbreviate this as \(\
ensuremath{\ensuremath{\mathsf{ax}}[n]}\).

In many contexts, there is only one axis with name \(\ensuremath{\
mathsf{ax}}\), and so we refer to the axis simply as \(\ensuremath{\
mathsf{ax}}\). The context always makes it clear whether \(\
ensuremath{\mathsf{ax}}\) is a name or an axis. If \(\ensuremath{\
mathsf{ax}}\) is an axis, we write \(\ind(\ensuremath{\mathsf{ax}})\)
for its index set, and we write \(|\ensuremath{\mathsf{ax}}|\) as
shorthand for \(|\ind(\ensuremath{\mathsf{ax}})|\).

A shape is a set of axes, written \(\ensuremath{\ensuremath{\mathsf
{ax_\text{$1$}}}[I_1]} \times \cdots \times \ensuremath{\ensuremath{\
mathsf{ax_\text{$r$}}}[I_r]}\), where \(\ensuremath{\mathsf{ax_\text
{$1$}}}, \ldots \ensuremath{\mathsf{ax_\text{$r$}}}\) are pairwise
distinct names. We write \(\emptyset\) for the empty shape. A shape
defines a set of records: \[\rec (\ensuremath{\ensuremath{\mathsf{ax_
\text{$1$}}}[I_1]} \times \cdots \times \ensuremath{\ensuremath{\
mathsf{ax_\text{$r$}}}[I_r]}) = \left\{\{\ensuremath{\ensuremath{\
mathsf{ax_\text{$1$}}}(i_1)}, \ldots, \ensuremath{\ensuremath{\mathsf
{ax_\text{$r$}}}(i_r)}\} \mid i_1 \in I_1, \ldots, i_r \in I_r\right
\}.\]

We say two shapes \(\mathcal{S}\) and \(\mathcal{T}\) are compatible
if whenever \(\ensuremath{\ensuremath{\mathsf{ax}}[I]} \in \mathcal
{S}\) and \(\ensuremath{\ensuremath{\mathsf{ax}}[J]} \in \mathcal{T}
\), then \(I = J\). We say that \(\mathcal{S}\) and \(\mathcal{T}\)
are orthogonal if there is no \(\ensuremath{\mathsf{ax}}\) such that 
\(\ensuremath{\ensuremath{\mathsf{ax}}[I]} \in \mathcal{S}\) and \(\
ensuremath{\ensuremath{\mathsf{ax}}[J]} \in \mathcal{T}\) for any \(I
\), \(J\).

If \(t \in \rec \mathcal{T}\) and \(\mathcal{S} \subseteq \mathcal{T}
\), then we write \(\mathopen{}\left.t\right|_{\mathcal{S}}\) for the
unique record in \(\rec \mathcal{S}\) such that \(\mathopen{}\left.t\
right|_{\mathcal{S}} \subseteq t\).

5.2 Named tensors

Let \(F\) be a field and let \(\mathcal{S}\) be a shape. Then a named
tensor over \(F\) with shape \(\mathcal{S}\) is a mapping from \(\
mathcal{S}\) to \(F\). We write the set of all named tensors with
shape \(\mathcal{S}\) as \(F^{\mathcal{S}}\).

We don't make any distinction between a scalar (an element of \(F\))
and a named tensor with empty shape (an element of \(F^\emptyset\)).

If \(A \in F^{\mathcal{S}}\), then we access an element of \(A\) by
applying it to a record \(s \in \rec \mathcal{S}\); but we write this
using the usual subscript notation: \(A_s\) rather than \(A(s)\). To
avoid clutter, in place of \(A_{\{\ensuremath{\ensuremath{\mathsf{ax_
\text{$1$}}}(x_1)}, \ldots, \ensuremath{\ensuremath{\mathsf{ax_\text
{$r$}}}(x_r)}\}}\), we usually write \(A_{\ensuremath{\ensuremath{\
mathsf{ax_\text{$1$}}}(x_1)}, \ldots, \ensuremath{\ensuremath{\mathsf
{ax_\text{$r$}}}(x_r)}}\). When a named tensor is an expression like 
\((A+B)\), we surround it with square brackets like this: \([A+B]_{\
ensuremath{\ensuremath{\mathsf{ax_\text{$1$}}}(x_1)}, \ldots, \
ensuremath{\ensuremath{\mathsf{ax_\text{$r$}}}(x_r)}}\).

We also allow partial indexing. If \(A\) is a tensor with shape \(\
mathcal{T}\) and \(s \in \rec \mathcal{S}\) where \(\mathcal{S} \
subseteq \mathcal{T}\), then we define \(A_s\) to be the named tensor
with shape \(\mathcal{T} \setminus \mathcal{S}\) such that, for any \
(t \in \rec (\mathcal{T} \setminus \mathcal{S})\), \[\begin{aligned}
\left[A_s\right]_t &= A_{s \cup t}.\end{aligned}\] (For the edge case
\(\mathcal{T} = \emptyset\), our definitions for indexing and partial
indexing coincide: one gives a scalar and the other gives a tensor
with empty shape, but we don't distinguish between the two.)

5.3 Named tensor operations

In SS2, we described several classes of functions that can be extended
to named tensors. Here, we define how to do this for general
functions.

Let \(f \colon F^{\mathcal{S}} \rightarrow G^{\mathcal{T}}\) be a
function from tensors to tensors. For any shape \(\mathcal{S'}\)
orthogonal to both \(\mathcal{S}\) and \(\mathcal{T}\), we can extend
\(f\) to: \[\begin{aligned} f \colon F^{\mathcal{S} \cup \mathcal
{S'}} &\rightarrow G^{\mathcal{T} \cup \mathcal{S'}} \\ [f(A)]_r &= f
(A_r) \qquad \text{for all $r \in \rec\mathcal{S'}$.}\end{aligned}\]

If \(f\) is a multary function, we can extend its arguments to larger
shapes, and we don't have to extend all the arguments with the same
names. We consider just the case of two arguments; three or more
arguments are analogous. Let \(f \colon F^{\mathcal{S}} \times G^{\
mathcal{T}} \rightarrow H^{\mathcal{U}}\) be a binary function from
tensors to tensors. For any shapes \(\mathcal{S'}\) and \(\mathcal
{T'}\) that are compatible with each other and orthogonal to \(\
mathcal{S}\) and \(\mathcal{T}\), respectively, and \(\mathcal{S'} \
cup \mathcal{T'}\) is orthogonal to \(\mathcal{U}\), we can extend \
(f\) to: \[\begin{aligned} f \colon F^{\mathcal{S} \cup \mathcal{S'}}
\times G^{\mathcal{T} \cup \mathcal{T'}} &\rightarrow H^{\mathcal{U}
\cup \mathcal{S'} \cup \mathcal{T'}} \\ [f(A,B)]_r &= f\left(A_{\
mathopen{}\left.r\right|_{\mathcal{S'}}},B_{\mathopen{}\left.r\right|
_{\mathcal{T'}}}\right) \qquad \text{for all $r \in \rec (\mathcal
{S'} \cup \mathcal{T'})$.}\end{aligned}\]

All of the tensor operations described in SS2.2 can be defined in this
way. For example, the contraction operator can be defined as: \[\
begin{aligned} \mathbin{\underset{\ensuremath{\mathsf{ax}}}{\odot}} \
colon F^{\ensuremath{\ensuremath{\mathsf{ax}}[n]}} \times F^{\
ensuremath{\ensuremath{\mathsf{ax}}[n]}} &\rightarrow F \\ A \mathbin
{\underset{\ensuremath{\mathsf{ax}}}{\odot}} B &= \sum_{i=1}^n A_{\
ensuremath{\ensuremath{\mathsf{ax}}(i)}} B_{\ensuremath{\ensuremath{\
mathsf{ax}}(i)}}.\end{aligned}\]

6 Differentiation

If \(f\) is a function from order-\(m\) tensors to order-\(n\)
tensors, the partial derivatives of \(f\) (evaluated on a tensor \(X
\)) form an order-\((m+n)\) tensor: \(m\) "input" axes for the
directions in which \(X\) could change and \(n\) "output" axes for
the change in \(f(X)\).

For example, the derivative of a function from vectors to vectors is
a matrix (the Jacobian). But using matrix notation, there are
conflicting conventions about whether the first axis is the input
axis ("denominator layout") or the output axis ("numerator layout").
The derivative of a function from vectors to matrices or matrices to
vectors cannot be represented as a matrix at all, so one must resort
to flattening the matrices into vectors.

With tensors, taking derivatives of higher-order tensors with respect
to higher-order tensors is not difficult (Laue, Mitterreiter, and
Giesen 2018). With named tensors, we get the additional advantage of
using names to distinguish input and output axes.

6.1 Definition

Let \(f \colon \mathbb{R}^\mathcal{S} \rightarrow \mathbb{R}^\mathcal
{T}\). The derivative of \(f\) (evaluated at \(X\)) has an input axis
for each axis in \(\mathcal{S}\) and an output axis for each axis in 
\(\mathcal{T}\), and they have to have distinct names. So if \(\
mathcal{S} = \ensuremath{\mathsf{ax_1}} \times \cdots \times \
ensuremath{\mathsf{ax_\text{$r$}}}\), then for each axis name \(\
ensuremath{\mathsf{ax_\text{$i$}}}\), let \(\ensuremath{\mathsf{ax_\
text{$i$}^*}}\) be a new axis name, not in \(\mathcal{T}\), and let \
(\mathcal{S}^* = \ensuremath{\mathsf{ax_1^*}} \times \cdots \times \
ensuremath{\mathsf{ax_\text{$r$}^*}}\). If \(s = \{\ensuremath{\
ensuremath{\mathsf{ax_1}}(i_1)}, \ldots, \ensuremath{\ensuremath{\
mathsf{ax_\text{$r$}}}(i_r)}\}\), let \(s^* = \{\ensuremath{\
ensuremath{\mathsf{ax_1^*}}(i_1)}, \ldots, \ensuremath{\ensuremath{\
mathsf{ax_\text{$r$}^*}}(i_r)}\}\).

Then the derivative of \(f\) at \(X\) is the tensor with shape \(\
mathcal{S}^* \times \mathcal{T}\) such that for all \(s \in \rec\
mathcal{S}\) and \(t \in \rec\mathcal{T}\), \[\left[\frac{\partial}{\
partial X}f(X) \right]_{s^*,t} = \frac{\partial}{\partial X_s} [f(X)]
_t.\]

We'll often make use of the following generalization of the identity
matrix: \[\begin{aligned} I_\mathcal{S} &\in \mathbb{R}^{\mathcal{S}^
* \times \mathcal{S}} \\ [I_\mathcal{S}]_{s^*, s} &= \begin{cases} 1
& \text{if $s^* = s$} \\ 0 & \text{otherwise.} \end{cases}\end
{aligned}\]

6.2 Rules

Now we give some rules for computing derivatives. Unless otherwise
indicated, \(X\) has shape \(\mathcal{S}\), and \(U\) and \(V\) are
dependent on \(x\) and have shapes \(\mathcal{U}\) and \(\mathcal{V}
\), respectively. \[\begin{aligned} \frac{\partial}{\partial X}X &=
I_\mathcal{S} \\ \frac{\partial}{\partial X}U &= 0 && \text{$U$ does
not depend on $X$} \\ \frac{\partial}{\partial X}f(U) &= f'(U) \odot
\frac{\partial}{\partial X}U && f \colon \mathbb{R}\rightarrow \
mathbb{R}\\ \frac{\partial}{\partial X}(U + V) &= \frac{\partial}{\
partial X}U + \frac{\partial}{\partial X}V \\ \frac{\partial}{\
partial X}\sum\limits_{\ensuremath{\mathsf{ax}}} U &= \sum\limits_{\
ensuremath{\mathsf{ax}}} \frac{\partial}{\partial X}U \\ \frac{\
partial}{\partial X}(U \odot V) &= \frac{\partial}{\partial X}U \odot
V + U \odot \frac{\partial}{\partial X}V \\ \frac{\partial}{\partial
X}(U \mathbin{\underset{\ensuremath{\mathsf{ax}}}{\odot}} V) &= \frac
{\partial}{\partial X}U \mathbin{\underset{\ensuremath{\mathsf{ax}}}
{\odot}} V + U \mathbin{\underset{\ensuremath{\mathsf{ax}}}{\odot}} \
frac{\partial}{\partial X}V \\ \frac{\partial}{\partial X}\frac{U}{V}
&= \frac{\frac{\partial}{\partial X}U \odot V - U \odot \frac{\
partial}{\partial X}V}{V^2} \\ \frac{\partial}{\partial X}U_r &= \
left[\frac{\partial}{\partial X}U\right]_r && r \in \rec \mathcal{R},
\mathcal{R} \subseteq \mathcal{U} \\ \frac{\partial}{\partial X}U_{{\
ensuremath{\mathsf{ax1}}\rightarrow\ensuremath{\mathsf{ax2}}}} &= \
left[\frac{\partial}{\partial X}U\right]_{{\ensuremath{\mathsf{ax1}}\
rightarrow\ensuremath{\mathsf{ax2}}}}\end{aligned}\]

The chain rule above is for elementwise operations. The general chain
rule looks like this for functions of one and two variables; three or
more variables are analogous. \[\begin{aligned} \frac{\partial}{\
partial X}f(U) &= \frac{\partial}{\partial U} f(U) \mathbin{\underset
{\ensuremath{\mathsf{\mathcal{U}^* \mid \mathcal{U}}}}{\odot}} \frac
{\partial}{\partial X}U \\ \frac{\partial}{\partial X}f(U, V) &= \
frac{\partial}{\partial U} f(U, V) \mathbin{\underset{\ensuremath{\
mathsf{\mathcal{U}^* \mid \mathcal{U}}}}{\odot}} \frac{\partial}{\
partial X}U + \frac{\partial}{\partial V} f(U, V) \mathbin{\underset
{\ensuremath{\mathsf{\mathcal{V}^* \mid \mathcal{V}}}}{\odot}} \frac
{\partial}{\partial X}V\end{aligned}\] where \(\mathbin{\underset{\
ensuremath{\mathsf{ax1|ax2}}}{\odot}}\) contracts \(\ensuremath{\
mathsf{ax1}}\) in the left operand with \(\ensuremath{\mathsf{ax2}}\)
in the right operand: \(A \mathbin{\underset{\ensuremath{\mathsf{ax1|
ax2}}}{\odot}} B = \sum_i A_{\ensuremath{\ensuremath{\mathsf{ax1}}
(i)}} \odot B_{\ensuremath{\ensuremath{\mathsf{ax2}}(i)}}\).

6.3 Examples

Here's an example using these rules to derive the Jacobian for
softmax: \[\begin{aligned} \frac{\partial}{\partial X}(\mathop{\
underset{\ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} X) &= \frac{\
partial}{\partial X}\frac{\exp X}{\sum\limits_{\ensuremath{\mathsf
{ax}}} \exp X} \\ &= \frac{\exp X \odot \frac{\partial}{\partial X}X
\odot \sum\limits_{\ensuremath{\mathsf{ax}}} \exp X - \exp X \odot \
sum\limits_{\ensuremath{\mathsf{ax}}} (\exp X \odot \frac{\partial}{\
partial X}X)}{(\sum\limits_{\ensuremath{\mathsf{ax}}} \exp X)^2} \\ &
= \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} X \
odot \left(\frac{\partial}{\partial X}X - \mathop{\underset{\
ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} X \mathbin{\underset{\
ensuremath{\mathsf{ax}}}{\odot}} \frac{\partial}{\partial X}X\right)
\\ &= \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{softmax}}}
X \odot \left(I_\mathcal{S} - \mathop{\underset{\ensuremath{\mathsf
{ax}}}{\mathrm{softmax}}} X \mathbin{\underset{\ensuremath{\mathsf
{ax}}}{\odot}} I_\mathcal{S}\right).\end{aligned}\] To derive the
backpropagation rule: \[\begin{aligned} \frac{\partial}{\partial X}f
(\mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} X) &=
f'(\mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} X)
\mathbin{\underset{\ensuremath{\mathsf{\mathcal{S}^*|\mathcal{S}}}}{\
odot}} \frac{\partial}{\partial X}\mathop{\underset{\ensuremath{\
mathsf{ax}}}{\mathrm{softmax}}} X \\ &= f'(\mathop{\underset{\
ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} X) \mathbin{\underset{\
ensuremath{\mathsf{\mathcal{S}^*|\mathcal{S}}}}{\odot}} \mathop{\
underset{\ensuremath{\mathsf{ax}}}{\mathrm{softmax}}} X \odot \left
(I_\mathcal{S} - \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm
{softmax}}} X \mathbin{\underset{\ensuremath{\mathsf{ax}}}{\odot}} I_
\mathcal{S}\right).\end{aligned}\]

As another example, here are Jacobian and backpropagation rule for
Conv1d: \[\begin{aligned} \frac{\partial}{\partial X} \text{Conv1d}
(X) &= [W \mathbin{\underset{\ensuremath{\mathsf{chans,kernel}}}{\
odot}} C \mathbin{\underset{\ensuremath{\mathsf{seq}}}{\odot}} I_\
mathcal{S}]_{{\ensuremath{\mathsf{out}}\rightarrow\ensuremath{\mathsf
{seq}}}} \\ &= [W_{{\ensuremath{\mathsf{chans}}\rightarrow\ensuremath
{\mathsf{chans^*}}}} \mathbin{\underset{\ensuremath{\mathsf{kernel}}}
{\odot}} C_{\ensuremath{\mathsf{seq}}\rightarrow\ensuremath{\mathsf
{seq^*}}}]_{\ensuremath{\mathsf{out}}\rightarrow\ensuremath{\mathsf
{seq}}} \\ \frac{\partial}{\partial X} f(\text{Conv1d}(X)) &= f'(\
text{Conv1d}(X)) \mathbin{\underset{\ensuremath{\mathsf{seq^*|seq}}}
{\odot}} [W_{{\ensuremath{\mathsf{chans}}\rightarrow\ensuremath{\
mathsf{chans^*}}}} \mathbin{\underset{\ensuremath{\mathsf{kernel}}}{\
odot}} C_{\ensuremath{\mathsf{seq}}\rightarrow\ensuremath{\mathsf{seq
^*}}}]_{\ensuremath{\mathsf{out}}\rightarrow\ensuremath{\mathsf
{seq}}} \\ &= f'(\text{Conv1d}(X)) \mathbin{\underset{\ensuremath{\
mathsf{seq^*|out}}}{\odot}} (W_{\ensuremath{\mathsf{chans}}\
rightarrow\ensuremath{\mathsf{chans^*}}} \mathbin{\underset{\
ensuremath{\mathsf{kernel}}}{\odot}} C_{\ensuremath{\mathsf{seq}}\
rightarrow\ensuremath{\mathsf{seq^*}}}) \\ &= W_{\ensuremath{\mathsf
{chans}}\rightarrow\ensuremath{\mathsf{chans^*}}} \mathbin{\underset
{\ensuremath{\mathsf{kernel}}}{\odot}} C_{{\ensuremath{\mathsf{seq}}\
rightarrow\ensuremath{\mathsf{seq^*}}}} \mathbin{\underset{\
ensuremath{\mathsf{out|seq^*}}}{\odot}} f'(\text{Conv1d}(X)).\end
{aligned}\]

6.4 Broadcasting

If \(f \colon \mathbb{R}^\mathcal{S} \rightarrow \mathbb{R}^\mathcal
{T}\), then recall that \(f\) can be extended to \(\mathbb{R}^{\
mathcal{S} \cup \mathcal{S^+}}\) where \(\mathcal{S}\) and \(\mathcal
{S^+}\) are orthogonal.

It's more convenient here to notate the derivative of \(f\) as \(Df\)
. If \(f\) has two arguments, its partial derivatives are \(D_1 f\)
and \(D_2 f\).

Although \(Df\) extends to \(\mathbb{R}^{\mathcal{S} \cup \mathcal{S^
+}}\) using the usual broadcasting rules, the extension of the
derivative is unfortunately not the derivative of the extension. To
avoid confusion, write \(f^+\) for the extension: \[\begin{aligned} f
^+ \colon \mathbb{R}^{\mathcal{S} \cup \mathcal{S^+}} &\rightarrow \
mathbb{R}^{\mathcal{T} \cup \mathcal{S^+}} \\ f^+(X)_r &= f(X_r).\end
{aligned}\] Then the derivative of \(f^+\) is: \[\begin{aligned} Df^+
\colon \mathbb{R}^{\mathcal{S}^* \cup \mathcal{{S^+}}^* \cup \mathcal
{T} \cup \mathcal{S^+}} &\rightarrow \mathbb{R}^{\mathcal{T} \cup \
mathcal{S^+}} \\ Df^+(X) &= Df(X) \odot I_{\mathcal{S}^+}.\end
{aligned}\]

Similarly, if \(f \colon \mathbb{R}^\mathcal{S} \times \mathbb{R}^\
mathcal{T} \rightarrow \mathbb{R}^\mathcal{U}\), we can extend \(f\)
to \(f^+ \colon \mathbb{R}^\mathcal{S \cup S^+} \times \mathbb{R}^\
mathcal{T \cup T^+} \rightarrow \mathbb{R}^\mathcal{U \cup S^+ \cup T
^+}\). Then \[\begin{aligned} D_1 f^+(X, Y) &= D_1 f(X, Y) \odot I_{\
mathcal{S}^+} \\ D_2 f^+(X, Y) &= D_2 f(X, Y) \odot I_{\mathcal{T}^
+}.\end{aligned}\]

7 Extensions

7.1 Index types

We have defined an axis as a pair \(\ensuremath{\ensuremath{\mathsf
{ax}}[I]}\), where \(\ensuremath{\mathsf{ax}}\) is a name and \(I\)
is a set, usually \([n]\) for some \(n\). In this section, we
consider some other possibilities for \(I\).

7.1.1 Non-integral types

The sets \(I\) don't have to contain integers. For example, if \(V\)
is the vocabulary of a natural language (\(V = \{ \ensuremath{\mathsf
{cat}}, \ensuremath{\mathsf{dog}}, \ldots \}\)), we could define a
matrix of word embeddings: \[\begin{aligned} E &\in \mathbb{R}^{\
ensuremath{\ensuremath{\mathsf{vocab}}[V]} \times \ensuremath{\
ensuremath{\mathsf{emb}}[d]}}.\end{aligned}\]

7.1.2 Integers with units

If \(\ensuremath{\mathsf{u}}\) is a symbol and \(n > 0\), define \
([n]\ensuremath{\mathsf{u}} = \{1\ensuremath{\mathsf{u}}, 2\
ensuremath{\mathsf{u}}, \ldots, n\ensuremath{\mathsf{u}}\}\). You
could think of \(\ensuremath{\mathsf{u}}\) as analogous to a physical
unit, like kilograms. The elements of \([n]\ensuremath{\mathsf{u}}\)
can be added and subtracted like integers (\(a\ensuremath{\mathsf{u}}
+ b\ensuremath{\mathsf{u}} = (a+b)\ensuremath{\mathsf{u}}\)) or
multiplied by unitless integers (\(c \cdot a\ensuremath{\mathsf{u}} =
(c \cdot a) \ensuremath{\mathsf{u}}\)), but numbers with different
units are different (\(a \ensuremath{\mathsf{u}} \neq a \ensuremath{\
mathsf{v}}\)).

Then the set \([n]\ensuremath{\mathsf{u}}\) could be used as an index
set, which would prevent the axis from being aligned with another
axis that uses different units. For example, if we want to define a
tensor representing an image, we might write \[A \in \mathbb{R}^{\
ensuremath{\ensuremath{\mathsf{height}}[[h]\ensuremath{\mathsf
{pixels}}]} \times \ensuremath{\ensuremath{\mathsf{width}}[[w]\
ensuremath{\mathsf{pixels}}]}}.\] If we have another tensor
representing a go board, we might write \[B \in \mathbb{R}^{\
ensuremath{\ensuremath{\mathsf{height}}[[n]\ensuremath{\mathsf
{points}}]} \times \ensuremath{\ensuremath{\mathsf{width}}[[n]\
ensuremath{\mathsf{points}}]}},\] and even if it happens that \(h = w
= n\), it would be incorrect to write \(A+B\) because the units do
not match.

7.1.3 Tuples of integers

An index set could also be \([m] \times [n]\), which would be a way
of sneaking ordered indices into named tensors, useful for matrix
operations. For example, instead of defining an \(\text{inv}\)
operator that takes two subscripts, we could write \[\begin{aligned}
A &\in \mathbb{R}^{\ensuremath{\ensuremath{\mathsf{ax}}[{m\times
n}]}} = \mathbb{R}^{\ensuremath{\ensuremath{\mathsf{ax}}[{[m]\times
[n]}]}} \\ B &= \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm
{inv}}} A.\end{aligned}\] We could also define an operator \(\circ\)
for matrix-matrix and matrix-vector multiplication: \[\begin{aligned}
c &\in \mathbb{R}^{\ensuremath{\ensuremath{\mathsf{ax}}[n]}} \\ D &=
A \mathbin{\underset{\ensuremath{\mathsf{ax}}}{\circ}} B \mathbin{\
underset{\ensuremath{\mathsf{ax}}}{\circ}} c.\end{aligned}\]

7.2 Indexing with a tensor of indices

Contributors: Tongfei Chen and Chu-Cheng Lin

NumPy defines two kinds of advanced (also known as fancy) indexing:
by integer arrays and by Boolean arrays. Here, we generalize indexing
by integer arrays to named tensors. That is, if \(A\) is a named
tensor with \(D\) indices and \(\iota^1, \ldots, \iota^D\) are named
tensors, called "indexers," what is \(A_{\iota^1, \ldots, \iota^D}\)?

Advanced indexing could be derived by taking a function \[\begin
{aligned} \mathop{\underset{\ensuremath{\mathsf{ax}}}{\mathrm
{index}}} &\colon F^{\ensuremath{\ensuremath{\mathsf{ax}}[I]}} \times
I \rightarrow F \\ \mathop{\underset{\ensuremath{\mathsf{ax}}}{\
mathrm{index}}}(A, i) &= A_{\ensuremath{\ensuremath{\mathsf{ax}}(i)}}
\end{aligned}\] and extending it to higher-order tensors in its
second argument according to the rules in SS5.3. But because that's
somewhat abstract, we give a more concrete definition below.

We first consider the case where all the indexers have the same shape
\(\mathcal{S}\): \[\begin{aligned} A &\in F^{\ensuremath{\ensuremath
{\mathsf{ax_\text{$1$}}}[I_1]} \times \cdots \times \ensuremath{\
ensuremath{\mathsf{ax_\text{$D$}}}[I_D]}} \\ \iota^d &\in I_d^{\
mathcal{S}} & d &= 1, \ldots, D.\end{aligned}\] Then \(A_{\iota^1, \
ldots, \iota^D}\) is the named tensor with shape \(\mathcal{S}\) such
that for any \(s \in \rec{\mathcal{S}}\), \[\begin{aligned} [A_{\iota
^1, \ldots, \iota^D}]_s &= A_{\iota^1_s, \ldots, \iota^D_s}.\end
{aligned}\] More generally, suppose the indexers have different but
compatible shapes: \[\begin{aligned} A &\in F^{\ensuremath{\
ensuremath{\mathsf{ax_\text{$1$}}}[I_1]} \times \cdots \times \
ensuremath{\ensuremath{\mathsf{ax_\text{$D$}}}[I_D]}} \\ \iota^d &\in
I_d^{\mathcal{S}_d} & d &= 1, \ldots, D,\end{aligned}\] where the \(\
mathcal{S}_d\) are pairwise compatible. Then \(A_{\iota^1, \ldots, \
iota^D}\) is the named tensor with shape \(\mathcal{S} = \bigcup_d \
mathcal{S}_d\) such that for any \(s \in \rec{\mathcal{S}}\), \[\
begin{aligned} [A_{\iota^1, \ldots, \iota^D}]_s &= A_{\iota^1_{\
mathopen{}\left.s\right|_{\mathcal{S}_1}}, \ldots, \iota^D_{\mathopen
{}\left.s\right|_{\mathcal{S}_D}}}.\end{aligned}\]

Let's consider a concrete example in natural language processing.
Consider a batch of sentences encoded as a sequence of word vectors,
that is, a tensor \(X \in \mathbb{R}^{\ensuremath{\ensuremath{\mathsf
{batch}}[B]} \times \ensuremath{\ensuremath{\mathsf{sent}}[N]} \times
\ensuremath{\ensuremath{\mathsf{emb}}[E]}}\). For each sentence, we
would like to take out the encodings of a particular span for each
sentence \(b \in [B]\) in the batch, resulting in a tensor \(Y \in \
mathbb{R}^{\ensuremath{\ensuremath{\mathsf{batch}}[B]} \times \
ensuremath{\ensuremath{\mathsf{span}}[M]} \times \ensuremath{\
ensuremath{\mathsf{emb}}[E]}}\).

We create a indexer for the \(\ensuremath{\mathsf{sent}}\) axis: \(\
iota \in [N]^{\ensuremath{\ensuremath{\mathsf{batch}}[B]} \times \
ensuremath{\ensuremath{\mathsf{span}}[M]}}\) that selects the desired
tokens. Also define the function \[\begin{aligned} \mathop{\underset
{\ensuremath{\mathsf{ax}}}{\mathrm{arange}}}(I) &\in I^{\ensuremath{\
ensuremath{\mathsf{ax}}[I]}} \\ \left[\mathop{\underset{\ensuremath{\
mathsf{ax}}}{\mathrm{arange}}}(I)\right]_{\ensuremath{\ensuremath{\
mathsf{ax}}(i)}} &= i\end{aligned}\] which generalizes the NumPy
function of the same name.

Then we can write \[Y = X_{\ensuremath{\ensuremath{\mathsf{batch}}(\
iota)}, \ensuremath{\ensuremath{\mathsf{sent}}(\mathop{\underset{\
ensuremath{\mathsf{sent}}}{\mathrm{arange}}}(n))}, \ensuremath{\
ensuremath{\mathsf{emb}}(\mathop{\underset{\ensuremath{\mathsf{emb}}}
{\mathrm{arange}}}(E))}}.\]

8 Alternatives

A very frequently asked question is why we haven't used index
notation as used in physics, and the Einstein summation convention in
particular. 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 of
the index, and if an index appears twice on one side and not on the
other, there is an implicit summation over that index. \[\begin
{aligned} \text{Attention} \colon \mathbb{R}^{n' \times d_k} \times \
mathbb{R}^{n \times d_k} \times \mathbb{R}^{n \times d_v} &\
rightarrow \mathbb{R}^{n' \times d_v} \\ \left[\text{Attention}(Q, K,
V)\right]_{i'k} &= \softmax_i \left( \frac{Q_{i'j} K_{ij}}{\sqrt
{d_k}} \right) V_{ik}.\end{aligned}\] Because \(i'\) and \(k\) appear
on both sides, the equation must hold over all values of these
indices. But because \(j\) and \(k\) occur twice on only the
right-hand side, they are both summed over. We'd have to define
exactly what the \(i\) under softmax means (\(i\) is bound inside the
softmax and free outside it), and since softmax doesn't distribute
over addition, we'd need to clarify that the summation over \(j\)
occurs inside the softmax.

Other than that, this is concise and unambiguous. But it doesn't
really solve the main problem we set out to solve, which is that
ordered axes force the author and reader to remember the purpose of
each axis. The indices do act as symbolic names for axes (indeed, in
abstract index notation, 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 a little bit more verbose than our
notation, particularly for reductions and contractions: \[\begin
{aligned} C &= \max_i A_i & C &=\mathop{\underset{\ensuremath{\mathsf
{ax}}}{\mathrm{max}}} A \\ C &= A_i B_i & C &= A \mathbin{\underset{\
ensuremath{\mathsf{ax}}}{\odot}} B.\end{aligned}\]

Finally, index notation requires us to write out all indices
explicitly. So if we wanted to extend attention to multiple heads and
minibatches, we would write: \[\begin{gathered} \text{Attention} \
colon \mathbb{R}^{B \times H \times n' \times d_k} \times \mathbb{R}^
{B \times H \times n \times d_k} \times \mathbb{R}^{B \times H \times
n \times d_v} \rightarrow \mathbb{R}^{B \times H \times n' \times
d_v} \\ \left[\text{Attention}(Q, K, V)\right]_{bhi'k} = \softmax_i \
left( \frac{Q_{bhi'j} K_{bhij}}{\sqrt{d_k}} \right) V_{bhik}.\end
{gathered}\] We could adopt a convention that extends a function on
tensors to tensors that have extra axes to the left, but such
conventions tend to lead to messy reordering and squeezing/
unsqueezing of axes. Named axes make this unnecessary.

Acknowledgements

Thanks to Ekin Akyurek, 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).

References

Chen, Tongfei. 2017. "Typesafe Abstractions for Tensor Operations."
In Proceedings of the 8th Acm Sigplan International Symposium on
Scala, 45-50. SCALA 2017. https://doi.org/10.1145/3136000.3136001.

Harris, Charles R., K. Jarrod Millman, Stefan J. van der Walt, Ralf
Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, et al. 2020.
"Array Programming with NumPy." Nature 585 (7825): 357-62. https://
doi.org/10.1038/s41586-020-2649-2.

Hoyer, Stephan, and Joe Hamman. 2017. "xarray: N-D Labeled Arrays and
Datasets in Python." Journal of Open Research Software 5 (1): 10.
https://doi.org/http://doi.org/10.5334/jors.148.

Laue, Soeren, Matthias Mitterreiter, and Joachim Giesen. 2018.
"Computing Higher Order Derivatives of Matrix and Tensor
Expressions." In Advances in Neural Information Processing Systems,
edited by S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N.
Cesa-Bianchi, and R. Garnett, 31:2750-9. Curran Associates, Inc.
https://proceedings.neurips.cc/paper/2018/file/
0a1bf96b7165e962e90cb14648c9462d-Paper.pdf.

Maclaurin, Dougal, Alexey Radul, Matthew J. Johnson, and Dimitrios
Vytiniotis. 2019. "Dex: Array Programming with Typed Indices." In
NeurIPS Workshop on Program Transformations for Ml. https://
openreview.net/forum?id=rJxd7vsWPS.

Paszke, Adam, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury,
Gregory Chanan, Trevor Killeen, et al. 2019. "PyTorch: An Imperative
Style, High-Performance Deep Learning Library." In Advances in Neural
Information Processing Systems 32, edited by H. Wallach, H.
Larochelle, A. Beygelzimer, F. d'Alche-Buc, E. Fox, and R. Garnett,
8024-35. Curran Associates, Inc. http://papers.neurips.cc/paper/
9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
.

Rush, Alexander. 2019. "Named Tensors." https://github.com/harvardnlp
/NamedTensor.

Sinha, Nishant. 2018. "Tensor Shape (Annotation) Library." https://
github.com/ofnote/tsalib.

Torch Contributors. 2019. "Named Tensors." https://pytorch.org/docs/
stable/named_tensor.html.

Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion
Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017.
"Attention Is All You Need." In Advances in Neural Information
Processing Systems, edited by I. Guyon, U. V. Luxburg, S. Bengio, H.
Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, 30:5998-6008.
Curran Associates, Inc. https://proceedings.neurips.cc/paper/2017/
file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.