matmul

Documentation

treetensor.torch.matmul(input, other, *args, **kwargs)[source]

In treetensor, you can create a matrix product with treetensor.torch.matmul().

Examples:

>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.matmul(
...     torch.tensor([[1, 2], [3, 4]]),
...     torch.tensor([[5, 6], [7, 2]]),
... )
tensor([[19, 10],
        [43, 26]])

>>> ttorch.matmul(
...     ttorch.tensor({
...         'a': [[1, 2], [3, 4]],
...         'b': {'x': [3, 4, 5, 6]},
...     }),
...     ttorch.tensor({
...         'a': [[5, 6], [7, 2]],
...         'b': {'x': [4, 3, 2, 1]},
...     }),
... )
<Tensor 0x7f2e74883f40>
├── a --> tensor([[19, 10],
│                 [43, 26]])
└── b --> <Tensor 0x7f2e74886430>
    └── x --> tensor(40)

Torch Version Related

This documentation is based on torch.matmul in torch v2.4.1+cu121. Its arguments’ arrangements depend on the version of pytorch you installed.

If some arguments listed here are not working properly, please check your pytorch’s version with the following command and find its documentation.

1
python -c 'import torch;print(torch.__version__)'

The arguments and keyword arguments supported in torch v2.4.1+cu121 is listed below.

Description From Torch v2.4.1+cu121

torch.matmul(input, other, *, out=None) → Tensor

Matrix product of two tensors.

The behavior depends on the dimensionality of the tensors as follows:

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

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

  • 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.

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

  • 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. 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. 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. The non-matrix (i.e. batch) dimensions are broadcasted (and thus must be broadcastable). For example, if input is a \((j \times 1 \times n \times n)\) tensor and other is a \((k \times n \times n)\) tensor, out will be a \((j \times k \times n \times n)\) tensor.

    Note that the broadcasting logic only looks at the batch dimensions when determining if the inputs are broadcastable, and not the matrix dimensions. For example, if input is a \((j \times 1 \times n \times m)\) tensor and other is a \((k \times m \times p)\) tensor, these inputs are valid for broadcasting even though the final two dimensions (i.e. the matrix dimensions) are different. out will be a \((j \times k \times n \times p)\) tensor.

This operator supports TensorFloat32.

Note

The 1-dimensional dot product version of this function does not support an out parameter.

Arguments:

input (Tensor): the first tensor to be multiplied other (Tensor): the second tensor to be multiplied

Keyword args:

out (Tensor, optional): the output tensor.

Example:

>>> # vector x vector
>>> tensor1 = torch.randn(3)
>>> tensor2 = torch.randn(3)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([])
>>> # matrix x vector
>>> tensor1 = torch.randn(3, 4)
>>> tensor2 = torch.randn(4)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([3])
>>> # batched matrix x broadcasted vector
>>> tensor1 = torch.randn(10, 3, 4)
>>> tensor2 = torch.randn(4)
>>> torch.matmul(tensor1, tensor2).size()
torch.Size([10, 3])
>>> # 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])
>>> # 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])