Where you put the brackets

Matrix multiplication is associative, so these are the same matrix:

(A @ B) @ C
A @ (B @ C)

They are not the same computation.

A (n, r)    B (r, n)    C (n, r)      with n large and r small

(A @ B) @ C     A @ B is (n, n)       n^2 elements
A @ (B @ C)     B @ C is (r, r)       r^2 elements

With n = 1000 and r = 2:

(A @ B) @ C     peak intermediate  1,000,000 elements
A @ (B @ C)     peak intermediate      2,000 elements

Five hundred times the memory, and the same number of matmuls. The only difference is where the brackets went.

This is not a micro-optimisation

It is the difference between running and not running. The (n, n) intermediate for a realistic sequence length is what makes naive attention quadratic in memory, and the whole family of linear-attention methods is fundamentally this observation: reassociate so the large axis never meets itself.

The same idea is behind LoRA. W + BA with B as (n, r) and A as (r, n) is never materialised as an (n, n) update; you apply B(Ax) to the activations instead and the big matrix never exists.

The rule

When chaining matmuls, multiply the small dimensions together first. Count the intermediate’s shape before you write the line: it is the product of the outer dimension of the left operand and the inner dimension of the right, and you can read it off without running anything.

Python’s @ is left-associative, so A @ B @ C means (A @ B) @ C and takes the expensive route by default. The brackets are not optional decoration.

Your task

def chain(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> torch.Tensor

a is (n, r), b is (r, n), c is (n, r). Return a @ b @ c without ever materialising an (n, n) intermediate.

The starter takes the default associativity.