oneflow.baddbmm

oneflow.baddbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None)Tensor

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.baddbmm.html.

Performs a batch matrix-matrix product of matrices in batch1 and batch2. input is added to the final result.

batch1 and batch2 must be 3-D tensors each containing the same number of matrices.

If batch1 is a \((b \times n \times m)\) tensor, batch2 is a \((b \times m \times p)\) tensor, then input must be broadcastable with a \((b \times n \times p)\) tensor and out will be a \((b \times n \times p)\) tensor.

\[\text{out}_i = \beta\ \text{input}_i + \alpha\ (\text{batch1}_i \mathbin{@} \text{batch2}_i)\]

If beta is 0, then input will be ignored, and nan and inf in it will not be propagated.

For inputs of type FloatTensor or DoubleTensor, arguments beta and alpha must be real numbers, otherwise they should be integers.

Args: input (Tensor): the tensor to be added batch1 (Tensor): the first batch of matrices to be multiplied batch2 (Tensor): the second batch of matrices to be multiplied

Keyword Arguments
  • beta (Number, optional) – multiplier for input (\(\beta\))

  • alpha (Number, optional) – multiplier for \(\text{{batch1}} \mathbin{{@}} \text{{batch2}}\) (\(\alpha\))

For example:

>>> import oneflow as flow
>>> input = flow.randn(10, 3, 5)
>>> batch1 = flow.randn(10, 3, 4)
>>> batch2 = flow.randn(10, 4, 5)
>>> of_out = flow.baddbmm(input, batch1, batch2)
>>> of_out.shape
oneflow.Size([10, 3, 5])