One weight matrix, every item

@ treats the last two axes as the matrix and every axis before them as batch. The batch axes broadcast by the ordinary rules; the matrix axes must agree in the ordinary matrix way.

(5, 2, 3) @ (3, 4)      ->  (5, 2, 4)
 ^^^^ batch  ^^^^ matrix, broadcast across the batch

(1, 2, 3) @ (5, 3, 4)   ->  (5, 3 is shared, so) (5, 2, 4)

So applying one weight matrix to a whole batch of activations needs nothing special:

x @ W       # x is (batch, time, in), W is (in, out) -> (batch, time, out)

No loop, no reshape, no repeat of W across the batch. This is the same right-alignment convention as the bias problem: features last, batch-like axes in front, and the library arranged so the common case costs nothing.

The two mistakes

Reshaping to force it:

x.reshape(-1, x.size(-1)) @ W      # works, and often copies

Fine when you also want the flattened shape, wasteful as a way of getting the matmul to happen at all.

Looping over the batch:

torch.stack([xi @ W for xi in x])

Correct, allocates the stack, and dispatches a matmul per item. On a batch of 256 that is 256 kernel launches instead of one.

What it actually dispatches

Worth looking at, because it is the previous track showing up inside the library:

x @ w   ->   aten.view, aten.mm, aten._unsafe_view      3 ops

PyTorch does not run a batched matmul at all here. It merges the leading axes into one, does a single 2-D mm, and views the result back. That is exactly pt-flatten-batch-dims: adjacent leading axes are mergeable as a view, so the batch can be folded away for free and handed to the one kernel that is most heavily optimised.

The loop cannot do that. On a batch of 3 it dispatches 7 ops against 3, and the gap grows with the batch.

The shape to check

If the result’s shape surprises you, print the two inputs’ shapes and mark off the last two of each. Everything else is a broadcast, and the broadcast rules are the ones from track 3.

Your task

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

x is (batch, time, in_features) and w is (in_features, out_features). Return (batch, time, out_features) in at most three dispatched operations.