mindspore.ops.bmm

mindspore.ops.bmm(input_x, mat2)[源代码]

基于batch维度的两个Tensor的矩阵乘法。

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

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

参数：

input_x (Tensor) - 输入相乘的第一个Tensor。其shape为 \((*B, N, C)\) ，其中 \(*B\) 表示批处理大小，可以是多维度， \(N\) 和 \(C\) 是最后两个维度的大小。
mat2 (Tensor) - 输入相乘的第二个Tensor。Tensor的shape为 \((*B, C, M)\) 。

返回：

Tensor，输出Tensor的shape为 \((*B, N, M)\) 。

异常：

ValueError - input_x 的维度小于 3 或者 mat2 的维度小于2。
ValueError - input_x 第三维的长度不等于 mat2 第二维的长度。

支持平台：

Ascend GPU CPU

样例：

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