Tensors and broadcasting

1. The idea

A tensor, in ML code, is an array with any number of axes: a number has 0, a vector 1, a matrix 2, and a batch of colour images 4: (batch, channels, height, width). Its shape lists the size of each axis, and many bugs are a shape you didn't expect.

Broadcasting is how NumPy and PyTorch combine arrays of different shapes without copying anything. Line the two shapes up from the right. Going axis by axis, the sizes must be equal, or one of them must be 1 (it is stretched to match), or one array has run out of axes (it counts as size 1). If any pair breaks the rule, you get an error.

This is what lets X − mean work on a whole batch: a (32, 10) batch minus a (10,) vector of feature means subtracts the right mean from every row. It is also the source of silent bugs: a (32,) vector of targets minus a (32, 1) column of predictions broadcasts to (32, 32), and the loss is wrong without any error.

2. Try it: line the shapes up

Pick a pair of shapes. The table applies the rule axis by axis from the right; for small 2-D cases the picture shows which numbers are real and which are stretched copies.

3. The math

4. The code (python)

import numpy as np

X = np.arange(12.0).reshape(3, 4)        # 3 examples, 4 features
mu = X.mean(axis=0)                      # one mean per feature, shape (4,)
print(mu.tolist())                       # [4.0, 5.0, 6.0, 7.0]
print((X - mu).shape)                    # (3, 4)
print((X - mu)[0].tolist())              # [-4.0, -4.0, -4.0, -4.0]

col = np.array([[0.0], [10.0], [20.0]])  # shape (3, 1)
row = np.array([1.0, 2.0, 3.0, 4.0])     # shape (4,)
print((col + row).shape)                 # (3, 4)

y = np.zeros(32)                         # targets, shape (32,)
y_hat = np.zeros((32, 1))                # predictions, shape (32, 1)
# a silent bug: no error, just the wrong shape
print((y - y_hat).shape)                 # (32, 32)

5. You’ll see this in

• Core Course chapter 6: arrays and batches
• Core Course chapter 16: convolutions B × C × H × W in practice
• Mean squared error, read with the seven questions where the (32,) − (32, 1) bug would bite

6. Check yourself