Broadcasting: arrays of different shapes

Read these first if they are new to you

Arrays in NumPy Intermediate

The idea in one paragraph

NumPy can add, subtract or multiply two arrays of different shapes. It lines up the two shapes from the right. Where one array has length 1 along an axis, or lacks that axis, NumPy repeats it to match the other. So subtracting one point, shape (3,), from a table of points, shape (4, 3), subtracts it from every row. This rule is broadcasting. NumPy does the repeating without making copies, so it costs no extra memory for the smaller array.

One row is repeated down every row of the table

Suppose you want to move four points so that their center sits at the origin. You subtract the center, (1, 2, 0.5), from each point. Without broadcasting you would write a loop over the rows. With broadcasting you write pts - center. NumPy treats the single row as if it were stacked four times.

(4, 3) minus (3,): the short array is repeated down the rows

Faint rows are not stored. NumPy reads the same row again for each row of the table. The result has shape (4, 3).

center.py

import numpy as np

pts = np.array([[0.0, 0.0, 0.0],
                [2.0, 0.0, 0.0],
                [2.0, 4.0, 0.0],
                [0.0, 4.0, 2.0]])
center = pts.mean(axis=0)
print(center.shape)
print(pts - center)        # (4, 3) minus (3,) gives (4, 3)
scale = np.array([[1.0], [2.0], [3.0], [4.0]])
print(scale.shape)
print(pts * scale)         # (4, 3) times (4, 1): row i times scale[i]

(3,)
[[-1.  -2.  -0.5]
 [ 1.  -2.  -0.5]
 [ 1.   2.  -0.5]
 [-1.   2.   1.5]]
(4, 1)
[[ 0.  0.  0.]
 [ 4.  0.  0.]
 [ 6. 12.  0.]
 [ 0. 16.  8.]]

The second case goes the other way. scale has one column, shape (4, 1), so NumPy repeats that column across the three columns. Each point is multiplied by its own number.

The rule: line the shapes up from the right, and stretch the 1s

Write the two shapes one above the other, aligned on the right.
If one shape is shorter, pad it on the left with 1s.
Compare the two lengths in each position. They fit if they are equal, or if one of them is 1.
The result takes the larger length in each position. If any position does not fit, NumPy raises an error.

A	B	B, padded	Result
`(4, 3)`	`(3,)`	`(1, 3)`	`(4, 3)`
`(4, 3)`	`(4, 1)`	`(4, 1)`	`(4, 3)`
`(3, 1, 2)`	`(1, 2, 2)`	`(1, 2, 2)`	`(3, 2, 2)`
`(4, 3)`	`(4,)`	`(1, 4)`	error: 3 and 4 do not fit

Try the rule: pick a shape for B

A has shape (4, 3). Before you choose, predict the result shape. Then find the one choice of B that fails.

B has shape

Solid cells are stored; dashed cells are repeats. A length-1 axis, or a missing one, is stretched. Any other mismatch is an error.

A length-1 axis turns two lists of points into a table of every pair

Broadcasting can also compare every point of one set with every point of another. Indexing with None inserts a new axis of length 1. Then a[:, None, :] has shape (3, 1, 2), and b[None, :, :] has shape (1, 2, 2). By the rule, their difference has shape (3, 2, 2): row i, column j holds a[i] − b[j].

Worked example: every distance between two small sets

Set a holds (0, 0), (3, 0) and (0, 4). Set b holds (0, 1) and (3, 4).

Shapes: a is (3, 2) and b is (2, 2). After adding axes they are (3, 1, 2) and (1, 2, 2), so the difference is (3, 2, 2).
Entry [0, 1] is a[0] − b[1] = (0 − 3, 0 − 4) = (−3, −4).
Square, add along the last axis, and take the root, as in the distance formula: √(9 + 16) = 5.
The full table of distances has shape (3, 2): one row per point of a, one column per point of b.

pairs.py

import numpy as np

a = np.array([[0.0, 0.0], [3.0, 0.0], [0.0, 4.0]])   # 3 points, shape (3, 2)
b = np.array([[0.0, 1.0], [3.0, 4.0]])               # 2 points, shape (2, 2)
diff = a[:, None, :] - b[None, :, :]                 # (3, 1, 2) - (1, 2, 2)
print(diff.shape)
d = np.sqrt((diff ** 2).sum(axis=-1))                # every distance, shape (3, 2)
print(d)

(3, 2, 2)
[[1.         5.        ]
 [3.16227766 4.        ]
 [3.         3.        ]]

When the shapes do not fit, NumPy says so:

import numpy as np

pts = np.zeros((4, 3))                     # a table of zeros, shape (4, 3)
# try runs the indented line. If that line fails with a ValueError,
# Python runs the lines under except instead, with the error in e.
try:
    pts - np.array([1.0, 2.0, 3.0, 4.0])
except ValueError as e:
    print(e)

operands could not be broadcast together with shapes (4,3) (4,)

Broadcasting can succeed when you meant it to fail

Subtract an array of shape (4,) from one of shape (4, 1), and you get shape (4, 4), not (4,). NumPy pairs every entry with every other entry, with no error. This is a common silent bug in code that computes a loss. Print the shapes when a result looks wrong. A table of every pair also grows fast: 100,000 points against 100,000 points in 3D is 3 × 10¹⁰ numbers, which takes 240 GB as 64-bit floats.

Check yourself

Give the result shape, or say it is an error: (5, 1) + (3,); (5, 3) + (5,); (2, 1, 3) − (4, 3).

Answer
(5, 3): pad (3,) to (1, 3), then both 1s stretch. Error: pad (5,) to (1, 5), and 3 does not fit 5. (2, 4, 3): pad (4, 3) to (1, 4, 3); the 1s stretch to 2 and 4.
You have 1000 points in an array of shape (1000, 3) and want every point moved up by 2 in z. Write the one-line array expression.

Answer
pts + np.array([0.0, 0.0, 2.0]). The shape (3,) row is repeated down all 1000 rows.
In the pairs example, what shape would a[:, None, :] - b[None, :, :] have if a held 100 points and b held 40, both in 3D?

Answer
(100, 40, 3). After the sum over the last axis, the distance table is (100, 40).

Go further

Optional background.

Broadcasting: General broadcasting rules NumPy documentation Article Intermediate

The official statement of the rule, with drawings of each case and a worked example of every pair.
A Visual Intro to NumPy: Matrix Arithmetic Jay Alammar Illustrated article Intermediate

Colored drawings of a table combined with a single row, the most common broadcast.

Where this shows up

No unit or explanation builds on this one yet.