Scale every row by its own weight

To multiply each row by its own weight, the weight axis must line up with the row axis. Broadcasting aligns dimensions from the right:

x:           (rows, cols)
w:                 (rows)

Here the length of w is compared with the column count, even though its values represent row weights. Matching sizes alone cannot express that meaning.

Insert an axis for the columns

w.unsqueeze(1) changes shape (rows,) to (rows, 1). Each weight then broadcasts across the columns of its own row:

import torch

x = torch.tensor([[1., 2.], [3., 4.]])
w = torch.tensor([10., 100.])
print((x * w).tolist())  # [[10.0, 200.0], [30.0, 400.0]]
print((x * w.unsqueeze(1)).tolist())  # [[10.0, 20.0], [300.0, 400.0]]

Both outputs have the same shape. The first scales columns, while the second performs the required row scaling.

x:                (rows, cols)
w.unsqueeze(1):   (rows,    1)
result:           (rows, cols)

When the wrong expression still runs

With unequal row and column counts greater than one, x * w raises a broadcasting error. Singleton dimensions are exceptions: with one column and several rows it broadcasts to (rows, rows), the wrong output shape. With a single row, its single weight happens to scale that row correctly.

Checking a rectangular example, a square example and singleton cases helps expose these different behaviors. Also check values: the wrong expression can have the expected shape.

Equivalent ways to align this vector

For the nonempty 1-D weights used here, these expressions all produce a storage-sharing column view:

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

unsqueeze names the axis insertion explicitly. It does not copy the weights; the multiplication produces the scaled output.

Match shapes from the right

Broadcasting lets an operation apply across compatible shapes without explicitly constructing repeated inputs. Compare dimensions from the right: each pair must be equal, or one must be 1. Missing leading dimensions behave like dimensions of 1. At each position, equal sizes stay the same. If one size is 1, the result uses the other size. This also covers empty dimensions: (0, 3) and (1, 3) broadcast to (0, 3), not (1, 3).

For example, (4, 3) plus (3,) applies the same three-feature bias to all four rows. Adding (4, 1) instead applies a different scalar to each row. A size-one axis lets the same value be reused along that axis.

This explains a subtle centring bug: reducing (B, D) over features gives (B,). Subtracting it directly aligns B with the feature axis. If B == D, the operation runs but uses the wrong means. Retaining the reduction axis as (B, 1) expresses the intended dependence.

Check how much memory the result needs

The input need not be physically copied, but the output can still be huge: (N, 1, D) minus (1, M, D) has shape (N, M, D). An in-place operation must fit its result in the existing destination shape. Explicit expansion may also create a view whose positions refer to the same storage, making independent writes unsafe.

Before implementing an operation, write the aligned shapes and the result shape. Then check which inputs are reused and how much space the output needs.