torch.bmm

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

Выполняет пакетное умножение матриц, хранящихся в input и mat2.

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

Если 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])

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

Параметры:

• input (Tensor) – первый пакет матриц для умножения
• mat2 (Tensor) – второй пакет матриц для умножения
• out_dtype (dtype) – тип данных выходного тензора. На CUDA и XPU для входных типов данных torch.float16/torch.bfloat16 поддерживается только torch.float32. Другие бэкенды (включая внешние ускорители) могут поддерживать дополнительные комбинации входных и выходных типов данных.

Именованные аргументы:

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