torch.mm vs torch.mul vs torch.matmul

Created using Gemini

Pytorch offeres three different functions to perform multiplication between two tensors. Lets understand how these functions are different from one another

Before we start a quick note on how to create row vector and column vector in Pytorch

# row vector
r = torch.tensor([1,2,3])
print(r.shape)
# >> torch.Size([3])

# column vector
c = r.T # transpose of r
# or
c = torch.tensor([[1],[2],[3]])
# >> torch.Size([3, 1])

torch.mm

Performs matrix multiplication of two tensors M1 and M2.

import torch
Output = torch.mm(M1,M2)

If M1 is a (n,m) tensor and M2 is a (m,p) tensor then Ouptut will be (n,p) tensor

• This function does not broadcast either M1/M2

#########################Examples of torch.mm#################################
# Matrix X Matrix
M1 = torch.rand(4,2)
M2 = torch.rand(2,3)
M1.mm(M2).shape
# >> torch.Size([4, 3])

# Matrix X Vector
M1 = torch.rand(4,2) #matrix
M2 = torch.rand(3) # vector
M1.mm(M2).shape
# >> RuntimeError: mat2 must be a matrix

# Vector X Vector
M1 = torch.rand(3)
M2 = torch.rand(3)
M1.mm(M2).shape
# >> RuntimeError: self must be a matrix

torch.matmul

Performs matrix product of two tensors with broadcasting with different behaviours depending on the tensor shapes as shown below.

• If both tensors are 1-dimensional, the dot product (scalar) is returned.

M1 = torch.rand(3)
M2 = torch.rand(3)
print(f"M1: {M1}")
print(f"M2: {M2}")
print(M1.matmul(M2))
print(f"M1.matmul(M2).shape: ",M1.matmul(M2).shape)

# >> M1: tensor([0.0211, 0.4164, 0.4223])
# >> M2: tensor([0.1626, 0.9487, 0.2893])
# >> tensor(0.5207)
# >> M1.matmul(M2).shape:  torch.Size([])

• If both arguments are 2-dimensional, the matrix-matrix product is returned.

M1 = torch.rand(3,2)
M2 = torch.rand(2,3)
print(f"M1.shape: {M1.shape}")
print(f"M2.shape: {M2.shape}")
# in this case matmul is same as mm
print((M1.matmul(M2)==M1.mm(M2)).all().sum()==1)
print(f"M1.matmul(M2).shape: ",M1.matmul(M2).shape)

# >> M1.shape: torch.Size([3, 2])
# >> M2.shape: torch.Size([2, 3])
# >> tensor(True)
# >> M1.matmul(M2).shape:  torch.Size([3, 3])

• If the first argument is 1-dimensional and the second argument is 2-dimensional, a 1 is prepended to its dimension for the purpose of the matrix multiply. After the matrix multiply, the prepended dimension is removed.

M1 = torch.rand(2)
M2 = torch.rand(2,3)
print(f"M1.shape: {M1.shape}")
print(f"M2.shape: {M2.shape}")
# M1 is converted from [2] -> [1,2]
# so [1,2] x [2,3] -> [1,3]
# we are comparing matmul with mm by manually reshaping M1 to [1,2] 
# as torch does it internally
print((M1.matmul(M2)==M1.view(1,2).mm(M2)).all().sum()==1)
print(f"M1.matmul(M2).shape: ",M1.matmul(M2).shape)

# >> M1.shape: torch.Size([2])
# >> M2.shape: torch.Size([2, 3])
# >> tensor(True)
# >> M1.matmul(M2).shape:  torch.Size([3])

• If the first argument is 2-dimensional and the second argument is 1-dimensional, the matrix-vector product is returned.

M1 = torch.tensor([[1,2,3,4]])
M2 = torch.tensor([2,2,2,2]) 
print(f"M1.shape: {M1.shape}")
print(f"M2.shape: {M2.shape}")
M1.matmul(M2)

# >> M1.shape: torch.Size([1, 4])
# >> M2.shape: torch.Size([4])
# >> tensor([20])

Batch matrix multiplication(BMM)

In order to understand next few scenarios we need to understand the concept of batched matrix multiplication. BMM involves matrices that are tensors with more than two dimensions. To grasp this process, we need to first comprehend PyTorch’s broadcasting mechanism.

Two tensors are “broadcastable” if the following rules hold:

• Each tensor has at least one dimension.
• When iterating over the dimension sizes, starting at the trailing dimension(the right most dim), the dimension sizes must either be equal, one of them is 1, or one of them does not exist.

Lets look at some examples to understand these rules

x=torch.empty(5,7,3)
y=torch.empty(5,7,3)
# same shapes are always broadcastable (i.e. the above rules always hold)

x=torch.empty((0,))
y=torch.empty(2,2)
# x and y are not broadcastable, because x does not have at least 1 dimension

# can line up trailing dimensions
x=torch.empty(5,3,4,1)
y=torch.empty(  3,1,1)
# x and y are broadcastable.
# 1st trailing dimension: both have size 1
# 2nd trailing dimension: y has size 1
# 3rd trailing dimension: x size == y size
# 4th trailing dimension: y dimension doesn't exist

x=torch.empty(5,2,4,1)
y=torch.empty(  3,1,1)
# x and y are not broadcastable, because in the 3rd trailing dimension 2 != 3

Now that we understand when broadcasting works lets see how BMM works.

If both arguments are at least 1-dimensional and at least one argument is N-dimensional (where N > 2), then a batched matrix multiply is returned.

• Case 1: If the second argument is 1-dimensional, a 1 is appended to its dimension for the purpose of the batched matrix multiple and removed after

# batched matrix x broadcasted vector
tensor1 = torch.randn(10, 3, 4)
tensor2 = torch.randn(4)
op = torch.matmul(tensor1, tensor2)

# second argument is tensor2 and its 1 dimensional
# so it will be appended with 1 that makes its shape [4,1]
# then we multiply each of [3,4] matrix from tensor1 with
# tensor2, [3,4] X [4,1] -> [3,1] => [3], as appended 1 dim is removed
# so there are 10 [3,4] matrices that will gives 10 of [3] vectors => [10,3]
op.shape

# >> torch.Size([10, 3])

• Case 2: If the first argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the batched matrix multiply and removed after.

# broadcasted vector x batched matrix 
tensor1 = torch.randn(4)
tensor2 = torch.randn(10, 4, 3)
op = torch.matmul(tensor1, tensor2)

# first argument is tensor1 and its 1 dimensional
# so it will be prepended with 1 that makes its shape [1,4]
# then we multiply each of [4,3] matrix from tensor2 with
# tensor1, [1,4]x[4,3] -> [1,3] => [3], as prepended 1 dim is removed
# so there are 10 [4,3] matrices that will gives 10 of [3] vectors => [10,3]
op.shape

# >> torch.Size([10, 3])

• The non-matrix (i.e. batch) dimensions are broadcasted (and thus must be broadcastable). For example, if input is a (j×1×n×n) tensor and other is a (k×n×n) tensor, out will be a (j×k×n×n) tensor.

# batched matrix x batched matrix
tensor1 = torch.randn(10, 3, 4)
tensor2 = torch.randn(10, 4, 5)
torch.matmul(tensor1, tensor2).size()

# >> torch.Size([10, 3, 5])
# because [3,4] x [4,5] -> [3,5] and there are 10 of them so its [10, 3,5]

# batched matrix x broadcasted matrix
tensor1 = torch.randn(10, 3, 4)
tensor2 = torch.randn(4, 5)
torch.matmul(tensor1, tensor2).size()

# >> torch.Size([10, 3, 5])

torch.mul

Performs an element-wise multiplication with broadcasting between inputs.

Case1: In the example shown below, A is a column vector with dimensions [4,1], and B is a row vector with dimensions [4]. In multiplication, the row vector B is transformed to match the dimensions of A, becoming [1,4]. When these two are multiplied, the result is a matrix with dimensions [4,4]

A = torch.tensor([[1],[2],[3],[4]])
B = torch.tensor([5,6,7,8])
print(f"A: {A.shape} x B: {B.shape} :",A.mul(B).shape)

# >> A: torch.Size([4, 1]) x B: torch.Size([4]) : torch.Size([4, 4])

Case2: In this one, C is a row vector with dimensions [1,4] and B is also a row vector with dimensions [4]. In multiplication, since both of them are row vectors it simply multiplies each element of C with each of element of B, that gives the output shape as [1,4]

A = torch.tensor([[1],[2],[3],[4]])
C = A.T
B = torch.tensor([5,6,7,8])
print(f"C: {C.shape} x B: {B.shape} :",C.mul(B).shape)

# >> C: torch.Size([1, 4]) x B: torch.Size([4]) : torch.Size([1, 4])

References:

[embed]torch.matmul - PyTorch 2.3 documentation Read the PyTorch Domains documentation to learn more about domain-specific librariespytorch.org

[embed]torch.mm - PyTorch 2.3 documentation Read the PyTorch Domains documentation to learn more about domain-specific librariespytorch.org

[embed]Broadcasting semantics - PyTorch 2.3 documentation Many PyTorch operations support NumPy's broadcasting semantics. See…pytorch.org

[embed]What's the difference between torch.mm, torch.matmul and torch.mul? After reading the pytorch documentation, I still require help in understanding the difference between torch.mm…stackoverflow.com.)