The outer product, three ways

Given a of length m and b of length n, produce the (m, n) matrix whose [i, j] entry is a[i] * b[j].

The broadcast

a[:, None] * b[None, :]

a[:, None]   (m, 1)
b[None, :]   (1, n)
------------------
           (m, n)

Both operands stretch: a across the columns, b down the rows. Every position ends up as a[i] * b[j], which is the definition.

This is the general pattern for any function of two vectors evaluated at every pair. Swap the * for a - and you have all pairwise differences; for == and you have a match matrix; for < and you have a causal mask:

torch.arange(t)[:, None] >= torch.arange(t)[None, :]     # lower triangle

That last line is the attention mask in every decoder you will read, and it is this problem with a comparison instead of a multiply.

The named version

torch.outer(a, b) exists and is clearer when the operation really is an outer product. Use it when it applies. Learn the broadcast anyway, because torch.outer only does multiplication and the pattern does everything.

The trap

Both of these are wrong and only one of them tells you:

a * b                # length mismatch, or a silent elementwise product when m == n
a[None, :] * b[:, None]     # the transpose of what you wanted

The second is the one to watch. It runs, it has the right shape when m == n, and it is b[i] * a[j].

Your task

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

Return the (m, n) outer product, in at most three dispatched operations, without torch.outer.

The starter loops.