mindspore.ops.bmm

mindspore.ops.bmm(input_x, mat2)[source]

Computes matrix multiplication between two tensors by batch.

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

The dim of input_x can not be less than 3 and the dim of mat2 can not be less than 2.

Parameters
  • input_x (Tensor) – The first tensor to be multiplied. The shape of the tensor is \((*B, N, C)\), where \(*B\) represents the batch size which can be multidimensional, \(N\) and \(C\) are the size of the last two dimensions.

  • mat2 (Tensor) – The second tensor to be multiplied. The shape of the tensor is \((*B, C, M)\).

Returns

Tensor, the shape of the output tensor is \((*B, N, M)\).

Raises
  • ValueError – If dim of input_x is less than 3 or dim of mat2 is less than 2.

  • ValueError – If the length of the third dim of input_x is not equal to the length of the second dim of mat2.

Supported Platforms:

Ascend GPU CPU

Examples

>>> input_x = Tensor(np.ones(shape=[2, 4, 1, 3]), mindspore.float32)
>>> mat2 = Tensor(np.ones(shape=[2, 4, 3, 4]), mindspore.float32)
>>> output = ops.bmm(input_x, mat2)
>>> print(output)
[[[[3. 3. 3. 3.]]
  [[3. 3. 3. 3.]]
  [[3. 3. 3. 3.]]
  [[3. 3. 3. 3.]]]
 [[[3. 3. 3. 3.]]
  [[3. 3. 3. 3.]]
  [[3. 3. 3. 3.]]
  [[3. 3. 3. 3.]]]]