Dense linear algebra
Apply a linear transformation represented by a weight matrix to a 2-D input matrix.
PyTorch · Function
torch.mm
torch.mm multiplies two 2-D tensors as matrices. Compatible shapes (n, m) and (m, p) produce a (n, p) tensor, without broadcasting.
How it works
torch.mm computes the matrix product of two 2-D tensors. If the first has shape (n, m) and the second (m, p), the result has shape (n, p); it does not add batch dimensions or broadcast inputs.
Syntax
torch.mm(input, mat2, *, out=None) → TensorParameters
| Name | Type | Description |
|---|---|---|
input | Tensor | First matrix to be multiplied. |
mat2 | Tensor | Second matrix to be multiplied. |
out | Tensor, keyword-only, optional | Optional destination tensor for the result. |
Returns
A Tensor containing the matrix product; for input shape (n, m) and mat2 shape (m, p), its shape is (n, p).
Key points
- The operation multiplies the first matrix by the second using standard matrix multiplication.
- On supported hardware, TensorFloat32 execution is available for this operator.
- An explicitly supplied output tensor determines the result layout.
- Without out, PyTorch derives the result layout from the first matrix.
Common mistakes
- Inputs must be ordinary compatible matrices; use torch.matmul when broadcasting across batch dimensions is required.
- Sparse inputs are supported, but automatic differentiation is documented with respect to strided inputs.
- Supplying out can affect the result layout, so avoid it when the destination layout is not intentional.
Example: Matrix multiplication
import torch
left = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
right = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
product = torch.mm(left, right)
print(product)
Multiplies a 2×3 matrix by a 3×2 matrix, producing a 2×2 tensor.
When to use it
Sparse matrix products
Multiply compatible sparse and dense 2-D tensors where the documented sparse layouts are supported.
Explicit 2-D products
Compute a fixed-size matrix product without batch broadcasting semantics.
Interview questions
QWhat does torch.mm compute, and what shapes must its inputs have?
torch.mm accepts two matrices and returns their matrix product. For shapes (n, m) and (m, p), the output shape is (n, p).
Follow-up: How would you handle batched matrix multiplication? Use torch.matmul, which is the documented alternative when broadcasting matrix products are needed.
QDoes torch.mm broadcast?
torch.mm does not broadcast its inputs. Both arguments are 2-D matrices whose inner dimensions must match.
Follow-up: What is the distinction from torch.matmul? torch.matmul supports broadcasting matrix products, while torch.mm is specifically the non-broadcasting matrix operation.
QWhat is the purpose of the out argument?
The optional out argument supplies the destination tensor. When it is supplied, its layout is used; otherwise, the result layout is inferred from the first input.
Follow-up: Why might omitting out be preferable? It lets PyTorch infer the result layout instead of imposing a possibly unintended destination layout.
Related
torch.matmultorch.bmmtorch.addmmtorch.sparse.mm