The Ultimate Guide to Matrix Multiplication with `torch.matmul()` in PyTorch

Matrix multiplication is a fundamental building block in various fields, including data science, computer graphics, and machine learning. PyTorch, a prominent machine learning library developed by Facebook, offers efficient ways to perform matrix multiplication using torch.matmul(). In this guide, we'll explore how to use torch.matmul() for practical applications, supported with code examples.

Understanding Matrix Multiplication

Before diving into torch.matmul(), it's crucial to understand the basics of matrix multiplication. Given two matrices, A and B, with A having dimensions (m x n) and B having dimensions (n x p), the result of multiplying A and B will be a new matrix C, with dimensions (m x p).

The mathematical formula for the elements of matrix C is:

C[i][j] = Σ (A[i][k] * B[k][j]) for all k from 1 to n

Using torch.matmul() for Matrix Multiplication

PyTorch’s torch.matmul() function provides a flexible and efficient way to perform matrix multiplication. Let’s explore its usage with some examples.

import torch

# Creating two matrices A and B
A = torch.tensor([[1, 2], [3, 4]])  # 2x2 matrix
B = torch.tensor([[5, 6], [7, 8]])  # 2x2 matrix

# Performing matrix multiplication
C = torch.matmul(A, B)

print(C)  # Output: tensor([[19, 22], [43, 50]])

Here, matrix C is computed by multiplying matrix A and B using torch.matmul().

Broadcasting and Batch Matrix Multiplication

One of the key features of torch.matmul() is its support for broadcasting. This means you can perform batch matrix multiplications efficiently when dealing with higher-dimensional tensors.

For instance, if you have a batch of matrices and want to perform matrix multiplication for each pair in the batch:

# Batch of matrices
A_batch = torch.rand(10, 2, 3)  # Batch of 10, 2x3 matrices
B_batch = torch.rand(10, 3, 4)  # Batch of 10, 3x4 matrices

# Perform batch matrix multiplication
C_batch = torch.matmul(A_batch, B_batch)

print(C_batch.size())  # Output: torch.Size([10, 2, 4])

In this example, torch.matmul() iteratively computes the product for each 2x3 matrix in A_batch and the corresponding 3x4 matrix in B_batch.

Handling Multidimensional Tensors

torch.matmul() can also handle tensors with more than two dimensions by employing matrix multiplication over the last two dimensions, while supporting broadcasting over the rest.

# Tensors with more than two dimensions
A_multi = torch.rand(2, 3, 4, 2)  # Shape: (2, 3, 4, 2)
B_multi = torch.rand(2, 3, 2, 5)  # Shape: (2, 3, 2, 5)

# Perform multidimensional matrix multiplication
C_multi = torch.matmul(A_multi, B_multi)

print(C_multi.size())  # Output: torch.Size([2, 3, 4, 5])

Here, matrix multiplication is performed over the last two dimensions due to PyTorch’s broadcasting rule, resulting in a 4x5 matrix for each subset defined by dimensions (2, 3).

Conclusion

The torch.matmul() function is a powerful tool in PyTorch for executing matrix multiplications effectively, supporting broadcasting and high-dimensional operations. Understanding how it works will offer clearer insights and significantly enhance any machine learning or data processing task, where matrix operations are frequent. As you delve deeper into PyTorch, mastering torch.matmul() is pivotal since it forms the backbone of many mathematical transformations.