Tensors, Axes, and Contractions

Extend matrix shape reasoning to arrays with any number of axes. Name axis meaning, distinguish tensor rank from matrix rank, compare dot and outer products, and track which indices remain or are summed away in a contraction.

A matrix has two axes, but many collections of data need more. A color image may have height, width, and channel axes. A batch of token representations may have batch, position, and feature axes. Such multi-axis arrays are commonly called tensors in numerical computing and machine learning. Their shapes become useful only when the meaning of every axis is known.

From Scalars to Multi-axis Arrays

The following objects form a useful progression:

Object	Number of axes	Example shape	Possible axis meanings
Scalar	0	$()$	one value
Vector	1	$(d,)$	feature
Matrix	2	$(B,d)$	example, feature
Three-axis tensor	3	$(B,L,d)$	example, position, feature
Four-axis tensor	4	$(B,H,W,C)$	example, height, width, channel

In this lesson, tensor rank means the number of axes. This usage differs from matrix rank, which measures the number of independent directions in a matrix and is studied later. A three-axis tensor has tensor rank three whether or not its entries contain redundant information.

The word tensor also has a more precise meaning in multilinear algebra and geometry. Numerical libraries often use it more broadly for a multi-axis array. That computational meaning is sufficient for the shape reasoning developed here.

Exercise: Read a tensor shape completely

A tensor $X$ has shape $(32,128,768)$ with axes $(\text{batch},\text{position},\text{feature})$ . Which statement is correct?

Choose the complete interpretation

X has tensor rank 3 and contains 768 features at each of 128 positions for each of 32 examples

X has tensor rank 768 because its feature axis has length 768

X has matrix rank 3 because its shape contains three numbers

Select one choice, then check.

HintCount axes before reading their lengths

The shape tuple contains three entries. Then attach batch, position, and feature meanings in that order.

SolutionDistinguish axis count from axis length

The shape contains three axes, so $X$ has tensor rank three in the computational sense. Axis $0$ selects one of 32 examples, axis $1$ selects one of 128 positions in that example, and axis $2$ selects one of 768 features at that position. The number $768$ is an axis length, not a rank.

Indices Identify Positions Along Named Axes

If:

X\in\mathbb{R}^{B\times L\times d},

then an entry can be written $X_{b\ell k}$ . The indices have separate roles:

$b$ selects an example from the batch;
$\ell$ selects a position within that example;
$k$ selects a feature at that position.

Fixing some indices produces a lower-dimensional slice. For example, $X_{b,:,:}$ is the $L\times d$ matrix for one example, while $X_{b,\ell,:}$ is the length- $d$ feature vector at one position. The colon notation means “keep every index along this axis.”

Axis order is part of the representation. Reordering $(B,L,d)$ to $(L,B,d)$ preserves the entries but changes which index position selects the batch. It must be represented by an explicit permutation or transpose operation; it is not merely a new verbal description of the same shape.

Outer Products Keep Both Input Axes

The dot product of two equal-length vectors contracts their shared coordinate index and produces a scalar:

\mathbf{a}^{\mathsf T}\mathbf{b} =\sum_i a_i b_i.

An outer product performs the pairwise multiplications without summing:

(\mathbf{a}\mathbf{b}^{\mathsf T})_{ij}=a_i b_j.

If $\mathbf{a}\in\mathbb{R}^m$ and $\mathbf{b}\in\mathbb{R}^n$ , then:

\mathbf{a}\mathbf{b}^{\mathsf T} \in\mathbb{R}^{m\times n}.

For example:

\begin{bmatrix} 2 \\ 3 \end{bmatrix} \begin{bmatrix} 4 & 5 & 6 \end{bmatrix} = \begin{bmatrix} 8 & 10 & 12 \\ 12 & 15 & 18 \end{bmatrix}.

The output retains one axis from each input. This is the opposite of summing the shared coordinate pairs into one value.

Exercise: Compare dot and outer products

Let $\mathbf{a},\mathbf{b}\in\mathbb{R}^4$ . What are the output shapes of $\mathbf{a}^{\mathsf T}\mathbf{b}$ and $\mathbf{a}\mathbf{b}^{\mathsf T}$ , respectively?

Choose the two output shapes

A scalar and a 4 × 4 matrix

A 4 × 4 matrix and a scalar

Two length-4 vectors

Select one choice, then check.

HintTrack the indices

The dot product contains $\sum_i a_i b_i$ . The outer product contains an unsummed pair $a_i b_j$ .

SolutionIdentify contracted and free indices

In the dot product, the shared index $i$ is summed away, leaving no free index and therefore a scalar. In the outer product, $i$ and $j$ remain free, giving one row and one column axis of length four. Its shape is $4\times4$ .

Contraction Sums Over Matched Axes

A tensor contraction multiplies entries and sums over one or more matched indices. Matrix multiplication is the familiar example:

C_{ij}=\sum_k A_{ik}B_{kj}.

The index $k$ is a contracted index: it occurs in both factors and is summed away. The indices $i$ and $j$ are free indices: they remain in the output and determine its axes.

This language extends to higher-rank arrays. Suppose a batch of position-wise feature vectors has shape:

X\in\mathbb{R}^{B\times L\times d_{\mathrm{in}}},

and a weight matrix has shape:

W\in\mathbb{R}^{d_{\mathrm{in}}\times d_{\mathrm{out}}}.

Applying $W$ to the feature vector at every batch and position gives:

Y_{b\ell j} =\sum_{k=1}^{d_{\mathrm{in}}}X_{b\ell k}W_{kj}.

Read this as: “for example $b$ , position $\ell$ , and output feature $j$ , multiply corresponding input features and weights across every input-feature index $k$ , then add those products.”

The input-feature index $k$ is contracted. The batch index $b$ , position index $\ell$ , and output-feature index $j$ remain. Therefore:

Y\in\mathbb{R}^{B\times L\times d_{\mathrm{out}}}.

Shape reasoning can be performed without expanding the entries:

(B,L,d_{\mathrm{in}}) \mathbin{@} (d_{\mathrm{in}},d_{\mathrm{out}}) \longrightarrow (B,L,d_{\mathrm{out}}).

Exercise: Find contracted and remaining axes

Let $X$ have shape $(16,50,64)$ with axes $(\text{batch},\text{position},\text{input feature})$ , and let $W$ have shape $(64,128)$ . For $Y_{b\ell j}=\sum_kX_{b\ell k}W_{kj}$ , which description is correct?

Choose the contraction and output shape

The input-feature index k is contracted, and Y has shape (16, 50, 128)

The batch index b is contracted, and Y has shape (50, 64, 128)

No index is contracted, and Y has shape (16, 50, 64, 128)

Select one choice, then check.

HintRead the index formula

The summation is over $k$ . The output is labelled by $b$ , $\ell$ , and $j$ .

SolutionRemove the summed index

The index $k$ selects the 64 input features in both $X$ and $W$ , so it is multiplied and summed away. The free indices have lengths $16$ , $50$ , and $128$ . Hence:

Y\in\mathbb{R}^{16\times50\times128}.

Every example and position remains, but its length- $64$ input feature vector is replaced by a length- $128$ output feature vector.

Broadcasting and Contraction Do Different Work

Broadcasting reuses values along compatible axes; it does not sum an axis away. Contraction multiplies values and explicitly sums over matched axes.

For example, if $X$ has shape $(B,L,d)$ and $\mathbf{b}$ has shape $(d,)$ , then $X+\mathbf{b}$ broadcasts the same length- $d$ vector across every batch and position. The output still has shape $(B,L,d)$ . By contrast, a dot product between the final axis of $X$ and a length- $d$ vector contracts that axis and produces shape $(B,L)$ .

Code can sometimes run even when the intended operation is wrong. Matching sizes are therefore only the first check. The operation must also match the meaning of the axes.

Verify Shapes and Values in NumPy

NumPy applies matrix multiplication to the final two axes while carrying earlier batch-like axes into the result:

Contract the feature axis of a three-axis tensor

Change the axis lengths or entries and compare the predicted and computed output shapes.

Ready to run.

The function np.einsum is named after the Einstein summation convention, a compact notation in which a repeated index is understood to be summed. Its first argument is a string that describes the input and output indices:

"blk,kj->blj"
  │   │    │
  X   W    Y

The part before -> describes the inputs in the same order as the arrays passed to the function:

blk assigns indices $b$ , $\ell$ , and $k$ to the three axes of X;
kj assigns indices $k$ and $j$ to the two axes of W.

The part after -> lists the output indices and their order. Here it keeps b, l, and j, producing an array with axes $(B,L,d_{\mathrm{out}})$ . The index k occurs in both inputs but not in the output, so einsum multiplies matching $k$ entries and sums over that axis:

Y_{b\ell j}=\sum_kX_{b\ell k}W_{kj}.

The letters have no built-in meaning to NumPy; b, l, and k are useful because we choose them to suggest batch, position, and feature. What matters to einsum is where each label appears. This notation can express many contractions and axis rearrangements, but ordinary @ is clearer for standard matrix multiplication. Use einsum when an explicit index pattern makes a higher-rank operation easier to verify, not merely because it is shorter.

A Reliable Shape-reading Method

For any tensor expression:

Name every input axis and write its length.
Mark axes that must match.
Identify indices that are summed away.
List the free indices in output order.
Check that the resulting axes still have the intended meaning.

This method applies to ordinary matrix products, batched neural network layers, image operations, and the attention calculations studied later. The chapter review now collects these tensor rules with the earlier matrix operations.