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])