mindspore.ops.bmm

查看源文件
mindspore.ops.bmm(input_x, mat2)[源代码]

对输入的两个tensor执行批量矩阵乘积。

\[\text{output}[..., :, :] = \text{matrix}(input_x[..., :, :]) * \text{matrix}(mat2[..., :, :])\]

input_x 的维度不能小于 3mat2 的维度不能小于 2

参数:
  • input_x (Tensor) - 输入的第一个Tensor。

  • mat2 (Tensor) - 输入的第二个Tensor。

返回:

Tensor

支持平台:

Ascend GPU CPU

样例:

>>> import mindspore
>>> input_x = mindspore.ops.arange(24, dtype=mindspore.float32).reshape(2, 4, 1, 3)
>>> mat2 = mindspore.ops.arange(72, dtype=mindspore.float32).reshape(2, 4, 3, 3)
>>> out = mindspore.ops.bmm(input_x, mat2)
>>> print(out)
[[[[  15.   18.   21.]]
  [[ 150.  162.  174.]]
  [[ 447.  468.  489.]]
  [[ 906.  936.  966.]]]
 [[[1527. 1566. 1605.]]
  [[2310. 2358. 2406.]]
  [[3255. 3312. 3369.]]
  [[4362. 4428. 4494.]]]]