One matrix, many batches

Use batch broadcasting when one matrix should be applied to several inputs. PyTorch offers matrix-multiplication functions with different shape requirements:

All require compatible matrix dimensions. matmul also has special rules for vectors; it does not accept scalar tensors.

One shared matrix

Here two inputs share the same weight matrix:

import torch

x = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]])  # (2, 1, 2)
w = torch.tensor([[[10.0], [1.0]]])              # (1, 2, 1)
print((x @ w).tolist())  # [[[12.0]], [[34.0]]]

The last two dimensions describe each matrix multiplication. The leading sizes, 2 and 1, are compatible for broadcasting, so the same matrix from w is used for both batch items.

torch.bmm(x, w) raises an error for these inputs because its batch sizes must match. When both batch sizes are one, it works; that special case can hide the starter's mistake.

Make the intended relationship clear

bmm is useful when you require two batches to pair item by item without broadcasting. matmul is useful when batch broadcasting is intended. Neither verifies what the axes mean in your application, so matching sizes alone do not prove the calculation is correct.

You could also explicitly expand w to the batch size before calling bmm. For this task, x @ w expresses the shared-matrix operation directly.

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.