Every pair of vector entries

An outer product multiplies every value in one vector by every value in another. For a of length m and b of length n, the result has shape (m, n) with entry [i, j] = a[i] * b[j].

Give each vector its own axis

import torch

a = torch.tensor([1., 2., 3.])
b = torch.tensor([10., 20.])
out = a[:, None] * b[None, :]
print(list(out.shape))  # [3, 2]
print(out.tolist())  # [[10.0, 20.0], [20.0, 40.0], [30.0, 60.0]]

a[:, None]: (m, 1)
b[None, :]: (1, n)
result:     (m, n)

The first vector varies down the rows and is reused across columns. The second varies across columns and is reused down the rows. Inserting these singleton axes creates views; multiplying them constructs the result matrix.

The pattern also works for comparisons

Replacing multiplication with subtraction gives all pairwise differences. A comparison gives a Boolean matrix. For example, comparing query positions with key positions creates a lower-triangular allowed-position mask:

positions = torch.arange(3)
allowed = positions[:, None] >= positions[None, :]
print(allowed.tolist())  # [[True, False, False], [True, True, False], [True, True, True]]

Here rows are queries, columns are keys, and True includes the current position and earlier positions. Check the receiving attention API's mask convention: some Boolean masks use True to exclude rather than allow. This is one way to construct a causal mask, not a requirement for every attention implementation.

Check the orientation

a * b does not form all pairs. It either multiplies aligned elements, broadcasts a singleton vector, or raises when lengths are incompatible. a[None, :] * b[:, None] forms all pairs in the reverse orientation, with shape (n, m). Even equal shapes do not establish correctness; compare values using different vectors.

torch.outer(a, b) is a direct API for the outer product of two vectors. For this exercise, use broadcasting instead to practice axis alignment.

Multiply rows by columns

For A with shape (M, K) and B with shape (K, N), the matrix product has shape (M, N). Each output entry is a dot product of one row of A and one column of B:

$$C_{ij}=\sum_{k=1}^{K}A_{ik}B_{kj}.$$

For example, multiplying row [1, 2] by column [3, 4] gives 1 * 3 + 2 * 4 = 11. Elementwise multiplication would instead give [3, 8]; it does not include the sum.

If examples are rows of X, with shape (B, D), and weights have shape (D, H), then XW produces H features for each of the B examples.

Keep batch axes separate from matrix axes

For batched matrix multiplication, the last two axes hold each matrix. Leading axes identify batches of matrices. A general matrix-multiply operation may broadcast compatible batch axes. A stricter batched operation may require them to match exactly. Use the operation specified by the exercise.

Einstein notation, often used with einsum, names the axes explicitly:

• bij,bjk->bik multiplies corresponding matrices in batch b and sums over j.
• bi,bj->bij forms an outer product for each example: output (b, i, j) is the product of input (b, i) and input (b, j).

An axis is summed out when its label is absent from the explicit output. A label that appears in both inputs is not necessarily summed: b stays in both outputs above.

Choose a multiplication order with small intermediate results

In exact arithmetic, (AB)C = A(BC), but the two orders can have very different costs. If the shapes are (1000, 10), (10, 10), and (10, 1), forming BC first needs a (10, 1) intermediate. Forming AB first needs a (1000, 10) intermediate and many more operations.

Floating-point rounding can also differ between orders, so do not expect every bit of the two results to match.