mindspore.ops.csr_mm

View Source On Gitee
mindspore.ops.csr_mm(a: CSRTensor, b: CSRTensor, trans_a: bool = False, trans_b: bool = False, adjoint_a: bool = False, adjoint_b: bool = False)[source]

Return the matrix multiplication result of the right-multiply matrix (dense or CSRTensor) of the CSRTensor. The CSRTensor with shape [M, N] needs to adapt the right matrix with shape [N, K] to get the dense matrix or CSRTensor with result [M, K].

Note

If right matrix is CSRTensor, currently only supports GPU backend. If right matrix is Tensor, currently supports CPU backend with LLVM no lower than 12.0.1, and GPU backend.

Parameters

  • a (CSRTensor) – Sparse CSR Tensor, rank should be 2.

  • b (CSRTensor) – Sparse CSR Tensor, rank should be 2.

  • trans_a (bool, optional) – whether to transpose CSRTensor a. Default: False .

  • trans_b (bool, optional) – whether to transpose CSRTensor b. Default: False .

  • adjoint_a (bool, optional) – whether to adjoint CSRTensor a. Default: False .

  • adjoint_b (bool, optional) – whether to adjoint CSRTensor b. Default: False .

Returns

CSRTensor.

Supported Platforms:

GPU

Examples

>>> from mindspore import Tensor, CSRTensor
>>> from mindspore import dtype as mstype
>>> from mindspore import ops
>>> a_shape = (4, 5)
>>> a_indptr = Tensor([0, 1, 1, 3, 4], dtype=mstype.int32)
>>> a_indices = Tensor([0, 3, 4, 0],dtype=mstype.int32)
>>> a_values = Tensor([1.0, 5.0, -1.0, -2.0], dtype=mstype.float32)
>>> b_shape = (5, 3)
>>> b_indptr = Tensor([0, 1, 1, 3, 3, 3], dtype=mstype.int32)
>>> b_indices = Tensor([0, 0, 1],dtype=mstype.int32)
>>> b_values = Tensor([2.0, 7.0, 8.0], dtype=mstype.float32)
>>> a = CSRTensor(a_indptr, a_indices, a_values, a_shape)
>>> b = CSRTensor(b_indptr, b_indices, b_values, b_shape)
>>> c = ops.csr_mm(a, b)
>>> print(c.shape)
(4, 3)
>>> print(c.values)
[2. -4.]
>>> print(c.indptr)
[0 1 1 1 2]
>>> print(c.indices)
[0 0]