Skip to content

← Matmul and einsum step 4 of 6

Easy Primitives

Star is not at

a * b      # elementwise, with broadcasting
a @ b      # matrix multiplication, contracting one axis

Different operations, different answers, and on square inputs the same output shape.

a = [[1, 2],     b = [[5, 6],
     [3, 4]]          [7, 8]]

a * b = [[ 5, 12],    a @ b = [[19, 22],
         [21, 32]]             [43, 50]]

Both (2, 2). Nothing raises. If your test data is square, and it usually is, * where you meant @ is a silent wrong answer that propagates.

This is the same failure shape as the missing keepdim from orientation, and it is the third time in this course that a square input has hidden a bug. The habit worth building is to make test shapes deliberately unequal.

Where it actually happens

Porting from NumPy code written before @ existed, where * on np.matrix meant matmul and on np.ndarray meant elementwise. Or translating a formula where juxtaposition means matrix product and someone typed the nearest operator.

The names

a @ b            torch.matmul(a, b)     broadcasting, any rank
                 torch.mm(a, b)         2-D only, no broadcasting
                 torch.bmm(a, b)        exactly 3-D, batch must match
a * b            torch.mul(a, b)        elementwise
                 (a * b).sum()          a dot product, the long way

The restricted names are useful precisely because they are restricted: mm refuses a 3-D input rather than silently treating it as a batch, so writing mm documents the rank you expect and enforces it.

Your task

def gram(x: torch.Tensor) -> torch.Tensor

x is (n, d). Return the (n, n) matrix of dot products between every pair of rows.

The starter uses *.