Broadcasting example
What is the shape of the resulting array:
res = einops.rearrange(np.arange(10), "b -> b 1") + np.arange(5)
print(res.shape)
Output: (10, 5)
Numpy broadcasting
Numpy's broadcasting has a match routine:
- Start from rightmost dimensions of both arrays.
- Aligned dimensions must be equal, or one of them must be 1.
- Continue until the shortest array is exhausted.
If there is a match, one or both arrays are stretched. Starting from the rightmost dimensions:
- A dimension of 1 is stretched (by repeating) to match the other dimension.
- if an array runs out of dimensions, the dimensions are padded on the left and then possibly stretched.
This behaviour is more intuitive in operations where both the expansion and stretching are confined to a single array:
# Add 0.5 to all elements in the array
img = np.arange(256*256*3).reshape(256, 256, 3)
res1 = img + 0.5
# Mask elements
assert mask.shape == (256, 256)
mask = einops.rearrange(mask, "h w -> h w 1")
res2 = img * mask
Stretching on both arrays
The example on the front side has the first array (shape (10,
1)) stretched to (10, 5) and the second array (shape
(5,)) expanded to (1, 5), then stretched to
(10, 5).
The takeaway is that array stretching can happen to both arrays as part of broadcasting.
Outer product, \( xx^{T} \)
We can calculate the outer product of two vectors using broadcasting:
x = einops.rearrange(np.arange(6), "n -> n 1")
xT = einops.rearrange(np.arange(6), "n -> 1 n")
res = x * xT
assert res.shape == (6, 6)
The rightmost dimension of x is stretched to n by
repeating, and the leftmost dimension of xT is stretched to
n by repeating. Then the two arrays are multiplied element-wise.