Tensors and Tensor Operations: Broadcasting, Reshaping and Einsum

In deep learning practice, most bugs are shape bugs. A batch of colour images is a 4-D array; a batch of token embeddings is 3-D; attention scores are 4-D. The mathematics of these objects is simple, but you must manipulate them fluently. This lecture is deliberately practical: by the end you should be able to read and write tensor code with confidence.

What is a tensor?#

In deep-learning libraries, a tensor is simply a multi-dimensional array with a data type and a device (CPU or GPU). Its rank (number of axes) and shape matter:

(Mathematicians use "tensor" for multilinear maps with transformation rules; in ML the word just means n-dimensional array.)

Indexing and slicing#

import numpy as np

x = np.arange(24).reshape(2, 3, 4)      # shape (2, 3, 4)
print(x[0].shape)          # (3, 4)   first item in batch
print(x[:, 1].shape)       # (2, 4)   second row of every item
print(x[..., -1].shape)    # (2, 3)   last element along the final axis
print(x[:, :, ::2].shape)  # (2, 3, 2) every other column
mask = x > 10
print(x[mask].shape)       # (13,)    boolean indexing flattens

Broadcasting#

Broadcasting lets operations combine tensors of different shapes without copying data. Rules (NumPy and PyTorch):

1. Align shapes from the right.
2. Two dimensions are compatible if they are equal or one of them is 1.
3. Missing leading dimensions are treated as 1.
4. Size-1 dimensions are virtually stretched to match.

X = np.random.rand(32, 10)        # batch of 32 feature vectors
mu = X.mean(axis=0)               # shape (10,)
Xc = X - mu                       # (32,10) - (10,) -> broadcast over batch

a = np.arange(3).reshape(3, 1)    # (3,1)
b = np.arange(4).reshape(1, 4)    # (1,4)
print((a + b).shape)              # (3,4): an "outer sum"

# Pairwise squared distances between rows of A (n,d) and B (m,d) -> (n,m)
A, B = np.random.rand(5, 3), np.random.rand(7, 3)
D = ((A[:, None, :] - B[None, :, :]) ** 2).sum(-1)
print(D.shape)                    # (5, 7)

Reshape, view, transpose and permute#

• reshape / view reinterpret the same elements in a new shape; the total size must match. Use -1 for one inferred dimension.
• transpose / permute reorder axes.

These are not interchangeable. Reshaping a (2, 3) matrix to (3, 2) is different from transposing it:

m = np.array([[1, 2, 3], [4, 5, 6]])
print(m.reshape(3, 2))   # [[1,2],[3,4],[5,6]]  - reads elements in memory order
print(m.T)               # [[1,4],[2,5],[3,6]]  - swaps axes

A common pattern in transformers — splitting the hidden dimension into attention heads:

import torch
B, T, H, Dh = 2, 5, 4, 8                  # batch, tokens, heads, head dim
x = torch.randn(B, T, H * Dh)             # (2, 5, 32)
heads = x.view(B, T, H, Dh).transpose(1, 2)   # (2, 4, 5, 8): batch, head, token, dim
merged = heads.transpose(1, 2).reshape(B, T, H * Dh)
print(torch.allclose(merged, x))          # True

In PyTorch, view requires contiguous memory; after transpose, use .reshape or call .contiguous() first.

Reductions#

Reductions collapse axes: sum, mean, max, argmax, logsumexp, norm. The axis (NumPy) or dim (PyTorch) argument specifies which axis disappears; keepdims=True keeps it as size 1, which is often what you need for subsequent broadcasting (e.g. normalising each row).

logits = np.random.randn(4, 10)
probs = np.exp(logits - logits.max(axis=1, keepdims=True))
probs /= probs.sum(axis=1, keepdims=True)   # each row sums to 1
print(probs.sum(axis=1))

Einstein summation#

einsum expresses products, transposes and reductions with index notation: indices that appear in inputs but not in the output are summed.

Q = torch.randn(2, 4, 5, 8)   # batch, heads, queries, dim
K = torch.randn(2, 4, 5, 8)   # batch, heads, keys, dim
scores = torch.einsum("bhqd,bhkd->bhqk", Q, K) / 8 ** 0.5
print(scores.shape)           # (2, 4, 5, 5)

Einsum makes the meaning of each axis explicit, which is why many researchers prefer it for complex operations. Libraries such as einops add readable rearrange operations, e.g. rearrange(x, 'b t (h d) -> b h t d', h=4).

Memory layout and devices#

Tensors live on a device; operations require all operands on the same device (x.to('cuda')). Moving data between CPU and GPU is slow — keep the training loop on the GPU. Choose data types deliberately (float32 for general training, bfloat16/float16 for mixed precision, int64 for class indices).