One matrix, many batches

Three functions do matrix multiplication and they differ in what they refuse.

torch.mm      exactly 2-D.  No batch, no broadcasting.
torch.bmm     exactly 3-D.  Batch sizes must be equal.
torch.matmul  any rank.     Batch axes broadcast. `@` is this.

So this raises:

a = torch.randn(1, 2, 3)
b = torch.randn(5, 3, 4)
torch.bmm(a, b)

RuntimeError: Expected size for first two dimensions of batch2
tensor to be [1, 3] but got [5, 3].

and this does not:

a @ b        # (5, 2, 4). The batch of 1 broadcast across the batch of 5.

The refusal is a feature

It is tempting to read bmm as “the worse matmul” and always use @. That gets it backwards. bmm refuses precisely the cases where broadcasting might not be what you meant, and a batch of 1 silently expanding to 5 is a bug about as often as it is an intention.

Choosing the restricted function is a way of writing down an assumption and having it checked. mm says “these are both plain matrices”; bmm says “these batches correspond one to one”. When that is true, saying so catches the day it stops being true.

This is the same argument as squeeze(dim=1) over squeeze() and sum(dim=0) over sum(): the form that names what it expects fails loudly when the expectation breaks.

When you do want the broadcast

Applying one shared matrix to every item of a batch, which is the previous problem. Then @ is right and the broadcast is the point.

Your task

def apply_shared(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor

x is (batch, n, d) and w is (1, d, m): a single matrix carrying a leading axis of 1. Return (batch, n, m), applying that one matrix to every item.

The starter uses bmm, which refuses to broadcast the batch of 1.