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