oneflow.bmm¶
-
oneflow.
bmm
()¶ Performs a batch matrix-matrix product of matrices stored in input and mat2.
input and mat2 must be 3-D tensors each containing the same number of matrices.
If input is a (b x n x m) tensor, mat2 is a (b x m x p) tensor, out will be a (b x n x p) tensor.
- Parameters
input (oneflow.Tensor) – the first batch of matrices to be multiplied
mat2 (oneflow.Tensor) – the second batch of matrices to be multiplied
For example:
>>> import oneflow as flow >>> import numpy as np >>> input1 = flow.randn(10, 3, 4) >>> input2 = flow.randn(10, 4, 5) >>> of_out = flow.bmm(input1, input2) >>> of_out.shape oneflow.Size([10, 3, 5])