Choose the order of matrix multiplications

The brackets in a chain of matrix products determine which intermediate matrix is built first. Consider:

a: (n, r)    b: (r, n)    c: (n, r)    with n much larger than r

Both (a @ b) @ c and a @ (b @ c) produce shape (n, r), but their first products have very different sizes.

The final output has 2,000 elements in both cases. These are individual tensor sizes, not peak process memory: inputs, simultaneously live tensors, backend workspaces and gradient tracking can also contribute.

The arithmetic changes too

Both expressions call matrix multiplication twice, but the matrix sizes determine the amount of work. Using conventional dense multiplication, the left grouping needs about 2 * n * n * r scalar multiplications; the right grouping needs about 2 * n * r * r.

For these shapes, computing b @ c first saves both intermediate storage and arithmetic. For other chains, compare the candidate shapes and costs rather than applying a rule based only on the smallest dimension. Python groups a @ b @ c from the left, so explicit brackets are needed for the smaller intermediate here.

Mathematical equality and rounding

Matrix multiplication is associative in exact arithmetic. Floating-point rounding can make the two groupings differ, so do not expect bit-for-bit equality for arbitrary tensors. This exercise uses small integer-valued inputs whose products and sums are exactly representable in float32.

Where the idea applies

For a low-rank update B @ A, applying it to a column vector as B @ (A @ x) avoids constructing the full update matrix. LoRA can use this fact during training; the update can also be merged into the base weights for inference.

Reassociation also helps some linear-attention formulations. It cannot simply move products through the softmax in ordinary attention: softmax(Q @ K.T) @ V is not generally equivalent to Q @ (K.T @ V).

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.