Elementwise or matrix multiplication?

* multiplies matching elements, with broadcasting where needed. @ performs matrix multiplication: each output entry sums products from one row of the left matrix and one column of the right.

import torch

a = torch.tensor([[1., 2.], [3., 4.]])
b = torch.tensor([[5., 6.], [7., 8.]])
print((a * b).tolist())  # [[5.0, 12.0], [21.0, 32.0]]
print((a @ b).tolist())  # [[19.0, 22.0], [43.0, 50.0]]

Both outputs have shape (2, 2). A shape check alone would not catch an accidental substitution. Test the values too, and include rectangular inputs to expose assumptions hidden by equal axis lengths.

Choose the operation

matmul requires both operands to have at least one dimension. For two vectors of equal length it returns their dot product. For two matrices, (n, d) @ (d, m) returns (n, m).

Compare every pair of rows

For a real matrix x with shape (n, d), a Gram matrix stores the dot product of each pair of rows. Entry [i, j] is:

sum(x[i, t] * x[j, t] for t in range(d))

Transposing x puts each original row into a column. Multiplying x @ x.t() then produces all row-pair dot products with shape (n, n). For strided tensors, the transpose shares storage; it does not copy the values. Creating a view and using its layout can still have a cost.

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.