🎓 A1Learner
🔍
0 of 79 lessons completed

PyTorch Broadcasting

Broadcasting stretches a smaller tensor inside an operation so its shape fits a bigger tensor. PyTorch does the stretching automatically, following fixed shape rules.

Adding a row the loop way

Adding one row to every row of a table takes a loop when each row is handled by hand.

table = torch.tensor([[1.0, 2.0, 3.0],
                      [4.0, 5.0, 6.0]])
row = torch.tensor([10.0, 20.0, 30.0])
result = torch.zeros_like(table)
for i in range(table.shape[0]):
    result[i] = table[i] + row
print(result)

The same job with broadcasting

Broadcasting stretches the row across the table. One line does the whole job, with the same result.

print(table + row)

A column stretches across

A column with shape (2, 1) stretches across the columns. Each row of the table gets its own added value.

col = torch.tensor([[100.0],
                    [200.0]])
print(table + col)

Example

import torch

table = torch.tensor([[2.0, 4.0, 6.0],
                      [4.0, 8.0, 10.0],
                      [6.0, 6.0, 14.0]])
means = table.mean(dim=0)
centred = table - means
print(means)
print(centred)

Key points

  • Broadcasting: a smaller tensor stretches to fit a bigger shape during an operation
  • Shapes line up from the right.
  • The last dimensions are compared first.
  • Two sizes fit when they are equal, or when one of them is 1
  • A size of 1 stretches to match the other size
  • Sizes that differ, with neither equal to 1, fail with an error
  • The stretched tensor is not stored.

In this example, table holds three rows of readings. table.mean(dim=0) gives the average down each column, so means is one row of three column averages. Subtracting means from table broadcasts that row across all three rows. Each cell ends up as its distance from its column average, so every column of centred averages zero.

Note

Broadcasting reuses the small tensor's values as it works.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A table has shape (3, 4). Which tensor can be added to it by broadcasting?