Skip to content

← Shapes and Strides step 5 of 9

Easy Primitives

Scale every row by its own weight

You met the silent broadcast in orientation. This is the tool for controlling it.

Broadcasting aligns shapes from the right and stretches any axis of length 1 to match. So the only lever you have is where the length-1 axes sit, and unsqueeze is how you put one where you need it.

The problem

x is (rows, cols). w is (rows,), one weight per row. Multiply each row by its weight.

x        (rows, cols)
w              (rows,)
aligned  (rows, cols)
                ^^^^
                w lines up with cols

Wrong axis. If rows == cols it runs and gives nonsense; otherwise it raises.

The fix

w.unsqueeze(1)      # (rows,) -> (rows, 1)
x                (rows, cols)
w.unsqueeze(1)   (rows,    1)
aligned          (rows, cols)
                          ^
                          stretches across the columns

Now each row’s weight applies to that whole row, which is what was meant.

Three spellings, one operation

w.unsqueeze(1)
w[:, None]
w.reshape(-1, 1)

All three produce (rows, 1) as a view over the same buffer. unsqueeze is the clearest about intent; [:, None] is the one you will meet most in other people’s code, borrowed from NumPy. Recognise all three.

A rule that survives contact

When you write a broadcast, write down both shapes and align them right. Two seconds of that catches nearly every broadcasting bug before it exists, and this is the habit the rest of the course assumes you have.

Your task

def scale_rows(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor

Multiply each row of x by the matching entry of w.