Subscripts instead of transposes

The attention scores from the previous problem, written twice:

torch.bmm(q, k.transpose(1, 2))
torch.einsum("bqd,bkd->bqk", q, k)

The second names every axis. b is batch in both inputs and in the output. d appears in both inputs and not in the output, so it is summed over. q and k appear once each and survive. That is the entire notation: repeated-and-dropped means contract, everything else is carried.

Nothing has to be transposed, because nothing is positional. The subscripts say which axis is which, so the axes can be in any order they like.

What it is actually good for

Reading. bqd,bkd->bqk tells you what the operation does; bmm after a transpose(1, 2) tells you what to type. Six months later that is a real difference, and it is the reason einsum is worth knowing even though it is rarely the fastest thing.

It also scales to operations that have no named function at all. Contracting three tensors, or two axes at once, is a subscript change rather than a sequence of reshapes.

What it is not good for, measured

In eager mode einsum is not cheaper. On a batched score computation:

torch.bmm(q, k.transpose(1, 2))        2 dispatched ops
torch.einsum("bqd,bkd->bqk", q, k)    12 dispatched ops

einsum parses the subscripts, plans a contraction order, and emits permutes and reshapes around a matmul. Under torch.compile that overhead disappears into the graph; in eager, in a hot inner loop, it is real.

So the honest rule is: einsum for clarity, and measure before putting it somewhere hot. A lot of writing about einsum implies it is faster. It is not, and the course would rather tell you than let you find out.

Your task

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

a is (batch, n, d) and b is (batch, m, d). Return (batch, n, m) where entry [x, i, j] is the dot product of a[x, i] and b[x, j].

Use torch.einsum. The starter has the subscripts wrong in a way that produces the right shape.