🎓 A1Learner
🔍

PyTorch Matrix Multiplication

Matrix multiplication multiplies each row of one tensor by each column of another and adds the products. PyTorch writes it with @ or torch.matmul.

Multiply by hand

Written out, matrix multiplication takes nested loops. Each output cell is one row times one column, added up.

a = torch.tensor([[1.0, 2.0],
                  [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0],
                  [7.0, 8.0]])
result = torch.zeros(2, 2)
for i in range(2):
    for j in range(2):
        result[i, j] = (a[i] * b[:, j]).sum()
print(result)

The @ operator

@ does the same work in one line. torch.matmul is the same operation written as a function.

print(a @ b)
print(torch.matmul(a, b))

When the inner dimensions clash

c has 3 columns and b has 2 rows. Running this raises a RuntimeError, and the message says the shapes cannot be multiplied.

c = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])
print(c @ b)

Example

import torch

prices = torch.tensor([[2.0],
                       [3.0],
                       [5.0]])
quantities = torch.tensor([[1.0, 2.0, 0.0],
                           [3.0, 0.0, 2.0]])
bills = quantities @ prices
print(bills)

Key points

  • Each output cell is one row times one column, added up
  • Inner dimensions: the columns of the first tensor and the rows of the second
  • The inner dimensions must match, or the multiplication fails with an error
  • The result takes the outer dimensions.
  • (2, 3) @ (3, 4) gives shape (2, 4).
  • torch.matmul: the function form of the @ operator

In this example, prices is a column holding the price of three items. quantities holds how many of each item two customers bought, one row per customer. The multiplication quantities @ prices pairs each customer row with the price column. Each result cell adds up one customer's items times their prices, so bills holds one total per customer.

Note

* multiplies cell by cell, and @ multiplies matrices.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What shape does torch.zeros(2, 3) @ torch.zeros(3, 5) have?