torch.bmm

torch.bmm(input, mat2, *, out=None) → Tensor

Выполняет произведение матриц в наборе, хранящихся в input и mat2.

input и mat2 должны быть 3-мерными тензорами, каждый из которых содержит одинаковое количество матриц.

Если input — тензор $(b \times n \times m)$, mat2 — тензор $(b \times m \times p)$, то out будет тензором $(b \times n \times p)$.

$$\text{out}_i = \text{input}_i \mathbin{@} \text{mat2}_i$$

Этот оператор поддерживает TensorFloat32.

На некоторых устройствах ROCm при использовании входных данных типа float16 этот модуль будет использовать разную точность для обратного распространения.

Примечание

Эта функция не выполняет векторизацию. Для векторизованного произведения матриц см. torch.matmul().

Параметры

• input (Tensor) – первый набор матриц для умножения
• mat2 (Tensor) – второй набор матриц для умножения

Аргументы ключевых слов

out (Tensor, необязательно) – тензор вывода.

Пример:

>>> input = torch.randn(10, 3, 4)
>>> mat2 = torch.randn(10, 4, 5)
>>> res = torch.bmm(input, mat2)
>>> res.size()
torch.Size([10, 3, 5])