Implementing Matrix Multiplication in PyTorch: An In-Depth Analysis from torch.dot to torch.matmul

Introduction

Matrix multiplication is a fundamental operation in tensor computations within the deep learning framework PyTorch. However, many developers transitioning from NumPy to PyTorch often encounter a common issue: when attempting to use the torch.dot function for matrix multiplication, it throws a "RuntimeError: 1D tensors expected, but got 2D and 2D tensors" error. This article aims to delve into the root causes of this phenomenon and systematically introduce the correct methods for implementing matrix multiplication in PyTorch.

Behavioral Differences Between torch.dot and np.dot

First, it is essential to understand the key distinctions between torch.dot and NumPy's np.dot function. In NumPy, np.dot is a versatile function: it computes the inner product for one-dimensional arrays and performs matrix multiplication for two-dimensional arrays. This flexibility allows the following code to work seamlessly:

import numpy as np
a = np.ones((3, 2))
b = np.ones((2, 1))
result = np.dot(a, b)  # Outputs a matrix of shape (3, 1)

In contrast, in PyTorch, torch.dot exhibits stricter behavior: it always treats input tensors as one-dimensional vectors and computes their dot product, regardless of the original dimensions of the tensors. Consequently, when trying to use torch.dot on two two-dimensional tensors, PyTorch flattens them into one-dimensional vectors. For example, a tensor a of shape (3, 2) is treated as a vector of length 6, and a tensor b of shape (2, 1) as a vector of length 2. Since the vector lengths do not match (6 vs. 2), the inner product cannot be computed, resulting in a runtime error. This design choice, discussed in PyTorch's GitHub issues, aims to maintain clear semantics for the function.

Correct Methods for Matrix Multiplication

To perform matrix multiplication in PyTorch, developers should use one of the following functions:

1. torch.mm: Specifically designed for multiplication of two-dimensional matrices and does not support broadcasting. For example:

   import torch
   a = torch.ones((3, 2))
   b = torch.ones((2, 1))
   result = torch.mm(a, b)  # Outputs a tensor of shape (3, 1)

2. torch.matmul: A more general function that supports matrix multiplication, vector dot products, and broadcasting operations. For two-dimensional inputs, its behavior is identical to torch.mm:

   result = torch.matmul(a, b)  # Outputs a tensor of shape (3, 1)

   Additionally, torch.matmul can handle dot products of one-dimensional vectors:

   a = torch.rand(n)
   b = torch.rand(n)
   result = torch.matmul(a, b)  # Outputs a scalar tensor

3. The @ operator in Python 3.5+: This serves as syntactic sugar for torch.matmul, making code more concise:

   result = a @ b

Advanced Features and Considerations

Beyond basic matrix multiplication, PyTorch offers several advanced capabilities:

• Broadcasting: torch.matmul supports broadcasting, allowing batch matrix operations. For instance, with three-dimensional tensors of shapes (b, m, k) and (b, k, j), it can perform batch matrix multiplication, whereas torch.mm only supports two-dimensional inputs.
• Element-wise Multiplication: For the Hadamard product (i.e., element-wise multiplication), the * operator can be used, provided the tensors have the same shape:

  c = a * b  # Element-wise multiplication

• Error Handling: Developers should be aware of function limitations. For example, torch.mm cannot be used with one-dimensional vectors or tensors of three or more dimensions, as it will throw an error.

Practical Recommendations

In practical applications, it is recommended to use torch.matmul or the @ operator, as they provide maximum flexibility while maintaining code readability. For performance-critical scenarios involving only two-dimensional matrices, torch.mm might offer slight advantages, but the differences are generally minimal. Always verify tensor shapes and refer to the official documentation for the latest behavior.

Conclusion

Through this analysis, we have learned that implementing matrix multiplication in PyTorch requires selecting the appropriate function based on specific needs. Avoiding direct use of torch.dot for matrix operations and instead adopting torch.mm, torch.matmul, or the @ operator ensures code correctness and efficiency. Combined with broadcasting and element-wise operations, PyTorch provides a robust toolkit for complex tensor manipulations.