mindspore 1.10.0__cp38-cp38-win_amd64.whl → 2.0.0rc1__cp38-cp38-win_amd64.whl

mindspore/ops/function/sparse_func.py

@@ -16,6 +16,7 @@
 """Defines sparse operators with functional form."""
 
 from __future__ import absolute_import
+from typing import Tuple
 from mindspore.ops.operations.sparse_ops import (
     DenseToCSRSparseMatrix,
     CSRSparseMatrixToSparseTensor,
@@ -23,14 +24,18 @@ from mindspore.ops.operations.sparse_ops import (
     SparseAdd,
     SparseMatrixAdd,
     SparseMatrixSoftmax,
+    SparseMatrixSparseMatMul,
     CSRSparseMatrixToDense
 )
+from mindspore import ops
+from mindspore.common import dtype as mstype
 from mindspore.ops.primitive import constexpr, Primitive
 from mindspore.ops.operations.array_ops import GatherNd, Coalesce
 from mindspore.ops.operations import _csr_ops
+from mindspore.ops import functional as F
 from mindspore.common import CSRTensor, COOTensor, Tensor
-from mindspore.common import dtype as mstype
-from mindspore.ops.composite.multitype_ops._constexpr_utils import raise_value_error, raise_type_error, make_tensor
+from mindspore.ops.composite.multitype_ops._constexpr_utils import raise_value_error, raise_type_error, make_tensor,\
+    promote_binary_dtype
 
 # utility functions and values
 gather_nd = GatherNd()
@@ -53,6 +58,20 @@ def _make_tensor(data):
     return Tensor(data)
 
 
+@constexpr
+def _make_tensor_with_dtype(data, dtype):
+    """Make Tensor with specific datatype"""
+    return Tensor(data, dtype=dtype)
+
+
+def _convert_shape(shape):
+    """Temporary solution to get shape value, will be removed when shape op is supported."""
+    if F.is_sequence_shape_unknown(shape):
+        return (-2,)
+    shape = [-1 if not F.isconstant(i) else i for i in shape]
+    return tuple(shape)
+
+
 def is_scalar(tensor):
     """Determine whether tensor input is a scalar tensor."""
     if tensor.size != 1:
@@ -60,7 +79,32 @@ def is_scalar(tensor):
     return len(tensor.shape) <= 2
 
 
-def coalesce(x_indices, x_values, x_shape):
+def promote_tensor(tensor_1, tensor_2):
+    """promote Tensor"""
+    dtype = promote_binary_dtype(tensor_1.dtype, tensor_2.dtype)
+    return tensor_1.astype(dtype), tensor_2.astype(dtype)
+
+
+def promote_csr(csr_tensor_1, csr_tensor_2):
+    """Type promotion for CSR tensor."""
+    indptr_1, indptr_2 = promote_tensor(csr_tensor_1.indptr, csr_tensor_2.indptr)
+    indices_1, indices_2 = promote_tensor(csr_tensor_1.indices, csr_tensor_2.indices)
+    values_1, values_2 = promote_tensor(csr_tensor_1.values, csr_tensor_2.values)
+    csr_tensor_1 = CSRTensor(indptr_1, indices_1, values_1, csr_tensor_1.shape)
+    csr_tensor_2 = CSRTensor(indptr_2, indices_2, values_2, csr_tensor_2.shape)
+    return csr_tensor_1, csr_tensor_2
+
+
+def promote_coo(coo_tensor_1, coo_tensor_2):
+    """promote COO-Tensor"""
+    indices_1, indices_2 = promote_tensor(coo_tensor_1.indices, coo_tensor_2.indices)
+    values_1, values_2 = promote_tensor(coo_tensor_1.values, coo_tensor_2.values)
+    coo_tensor_1 = COOTensor(indices_1, values_1, coo_tensor_1.shape)
+    coo_tensor_2 = COOTensor(indices_2, values_2, coo_tensor_2.shape)
+    return coo_tensor_1, coo_tensor_2
+
+
+def coalesce(x_indices: Tensor, x_values: Tensor, x_shape: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
     """
     Returns the coalesced sparse tensor of the input.
 
@@ -92,7 +136,7 @@ def coalesce(x_indices, x_values, x_shape):
         ValueError: If any of the values of elements of `x_indices` exceed the limit set by `x_shape`.
 
     Supported Platforms:
-        ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore
@@ -101,7 +145,7 @@ def coalesce(x_indices, x_values, x_shape):
         >>> x_indices = Tensor([[0, 0, 1], [1, 1, 2]], dtype=ms.int64)
         >>> x_values = Tensor([1, 5, 4], dtype=ms.float32)
         >>> x_shape = Tensor([3, 3], dtype=ms.int64)
-        >>> y_indices, y_values, y_shape = ops.coalesce(x_indices, x_values, x_shape)
+        >>> y_indices, y_values, y_shape = ops.Coalesce()(x_indices, x_values, x_shape)
         >>> print(y_indices)
         [[0 1]
          [1 2]]
@@ -122,7 +166,7 @@ coo_tensor_get_indices = Primitive('COOTensorGetIndices')
 coo_tensor_get_values = Primitive('COOTensorGetValues')
 
 
-def csr_div(x, y):
+def csr_div(x: CSRTensor, y: Tensor) -> Tensor:
     """
     Returns x / y where x is CSRTensor and y is Tensor.
 
@@ -148,28 +192,24 @@ def csr_div(x, y):
         if y.ndim > x.ndim:
             raise_value_error("dense tensor cannot broadcast to the sparse tensor.")
         return (x.values / y).reshape(x.values.shape)
-    return _csr_ops.CSRDiv()(x.indptr, x.indices, x.values, x.shape, y)
+    x_values, y = promote_tensor(x.values, y)
+    res_values = _csr_ops.CSRDiv()(x.indptr, x.indices, x_values, x.shape, y)
+    return CSRTensor(x.indptr, x.indices, res_values, x.shape)
 
 
 csr_gather = _csr_ops.CSRGather()
 
 
-def csr_mul(x, y):
+def csr_mul(x: CSRTensor, y: Tensor) -> CSRTensor:
     """
     Returns x * y where x is CSRTensor and y is Tensor.
 
-    Note:
-        This function returns the results of dense Tensor, represents the non-zero
-        values of the CSRTensor. If user expects a CSRTensor as output, please directly
-        use `*` operator instead. Only support dense tensor broadcast to sparse tensor
-        at the moment.
-
     Args:
         x (CSRTensor): Sparse CSR Tensor.
         y (Tensor): Dense Tensor, its shape must be able to broadcast to x.
 
     Returns:
-        Dense Tensor, represents the non-zero values of the result.
+        CSRTensor.
 
     Supported Platforms:
         ``GPU`` ``CPU``
@@ -180,10 +220,12 @@ def csr_mul(x, y):
         if y.ndim > x.ndim:
             raise_value_error("dense tensor cannot broadcast to the sparse tensor.")
         return (x.values * y).reshape(x.values.shape)
-    return _csr_ops.CSRMul()(x.indptr, x.indices, x.values, x.shape, y)
+    x_values, y = promote_tensor(x.values, y)
+    res_values = _csr_ops.CSRMul()(x.indptr, x.indices, x_values, x.shape, y)
+    return CSRTensor(x.indptr, x.indices, res_values, x.shape)
 
 
-def csr_mv(csr_tensor, dense):
+def csr_mv(csr_tensor: CSRTensor, dense: Tensor) -> Tensor:
     """
     Sparse matrix-vector multiplication.
 
@@ -197,10 +239,87 @@ def csr_mv(csr_tensor, dense):
     Supported Platforms:
         ``GPU`` ``CPU``
     """
-    return _csr_ops.CSRMV()(csr_tensor.indptr, csr_tensor.indices, csr_tensor.values, csr_tensor.shape, dense)
+    csr_tensor_values, dense = promote_tensor(csr_tensor.values, dense)
+    return _csr_ops.CSRMV()(csr_tensor.indptr, csr_tensor.indices, csr_tensor_values, csr_tensor.shape, dense)
 
 
-def csr_reduce_sum(csr_tensor, axis):
+def csr_mm(a: CSRTensor, b: CSRTensor, trans_a: bool = False, trans_b: bool = False,
+           adjoint_a: bool = False, adjoint_b: bool = False):
+    """
+    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:
+        Currently supports GPU backend with right matrix is CSRTensor.
+
+    Args:
+        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
+        >>> import mindspore.ops as 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]
+    """
+    if not isinstance(a, CSRTensor) or not isinstance(b, CSRTensor):
+        raise_type_error("For functional operator csr_mm, inputs a and b must be type of CSRTensor currently.")
+    a_batch_pointers = make_tensor([0, a.values.shape[0]], a.indices.dtype)
+    b_batch_pointers = make_tensor([0, b.values.shape[0]], b.indices.dtype)
+    a_shape = make_tensor(a.shape, a.indices.dtype)
+    b_shape = make_tensor(b.shape, b.indices.dtype)
+    sparse_matrix_sparse_matmul = SparseMatrixSparseMatMul(transpose_a=trans_a,
+                                                           transpose_b=trans_b,
+                                                           adjoint_a=adjoint_a,
+                                                           adjoint_b=adjoint_b)
+
+    _, _, c_indptr, c_indices, c_values = sparse_matrix_sparse_matmul(a_shape,
+                                                                      a_batch_pointers,
+                                                                      a.indptr,
+                                                                      a.indices,
+                                                                      a.values,
+                                                                      b_shape,
+                                                                      b_batch_pointers,
+                                                                      b.indptr,
+                                                                      b.indices,
+                                                                      b.values)
+    if trans_a or adjoint_a:
+        return CSRTensor(c_indptr, c_indices, c_values, (a.shape[1], b.shape[1]))
+    if trans_b or adjoint_b:
+        return CSRTensor(c_indptr, c_indices, c_values, (a.shape[0], b.shape[0]))
+    return CSRTensor(c_indptr, c_indices, c_values, (a.shape[0], b.shape[1]))
+
+
+def csr_reduce_sum(csr_tensor: CSRTensor, axis: int) -> Tensor:
     """
     Reduces a dimension of a CSRTensor by summing all elements in the dimension.
 
@@ -217,7 +336,7 @@ def csr_reduce_sum(csr_tensor, axis):
     return _csr_ops.CSRReduceSum()(csr_tensor.indptr, csr_tensor.indices, csr_tensor.values, csr_tensor.shape, axis)
 
 
-def csr_to_coo(tensor):
+def csr_to_coo(tensor: CSRTensor) -> COOTensor:
     """
     Converts a CSRTensor to COOTensor.
 
@@ -251,15 +370,15 @@ def csr_to_coo(tensor):
     """
     if not isinstance(tensor, CSRTensor):
         raise_type_error("For functional operator csr_to_coo, input argument must be a CSRTensor.")
-    if len(tensor.shape) != 2:
+    if len(_convert_shape(tensor.shape)) > 2:
         raise_value_error("Currently only support 2-D CSRTensor when converting to COOTensor.")
     shape = tensor.shape
     indices, values, _ = csr_sparse_matrix_to_sparse_tensor(Tensor(shape, mstype.int32), batch_csr_pointers_empty,
                                                             tensor.indptr, tensor.indices, tensor.values)
-    return COOTensor(indices, values, shape)
+    return COOTensor(indices, values, _convert_shape(shape))
 
 
-def csr_to_dense(csr_tensor):
+def csr_to_dense(csr_tensor: CSRTensor) -> Tensor:
     """
     Converts a CSRTensor to its dense form.
 
@@ -291,11 +410,24 @@ def csr_to_dense(csr_tensor):
     """
     if not isinstance(csr_tensor, CSRTensor):
         raise_type_error("For functional operator csr_to_dense, input argument must be a CSRTensor.")
-    if len(csr_tensor.shape) != 2:
+    if len(csr_tensor.shape) > 2:
         raise_value_error("Currently only support 2-D CSRTensor when converting to COOTensor.")
-    shape = csr_tensor.shape
-    return csr_sparse_matrix_to_dense(Tensor(shape, dtype=mstype.int32), batch_csr_pointers_empty,
-                                      csr_tensor.indptr, csr_tensor.indices, csr_tensor.values)
+
+    shape = _convert_shape(csr_tensor.shape)
+    dense_shape = Tensor(shape, dtype=mstype.int32)
+    batch_pointers = ops.concat((make_tensor([0]), ops.TensorShape()(csr_tensor.values))).astype("int32")
+    row_pointers = csr_tensor.indptr
+    col_indices = csr_tensor.indices
+    values = csr_tensor.values
+
+    valid_indices_dtype = [mstype.int32, mstype.int64]
+    if row_pointers.dtype in valid_indices_dtype and col_indices.dtype in valid_indices_dtype:
+        if row_pointers.dtype == mstype.int64 or col_indices.dtype == mstype.int64:
+            return csr_sparse_matrix_to_dense(dense_shape.astype(mstype.int64), batch_pointers.astype(mstype.int64),
+                                              row_pointers.astype(mstype.int64), col_indices.astype(mstype.int64),
+                                              values)
+
+    return csr_sparse_matrix_to_dense(dense_shape, batch_pointers, row_pointers, col_indices, values)
 
 
 # deprecated, will be removed once `csr_to_coo` supports all backends.
@@ -310,7 +442,7 @@ csr_tensor_get_indptr = Primitive('CSRTensorGetIndptr')
 csr_tensor_get_values = Primitive('CSRTensorGetValues')
 
 
-def dense_to_sparse_coo(tensor):
+def dense_to_sparse_coo(tensor: Tensor) -> COOTensor:
     """
     Convert a Tensor to COOTensor.
 
@@ -349,14 +481,14 @@ def dense_to_sparse_coo(tensor):
     """
     if not isinstance(tensor, Tensor):
         raise_type_error("For functional operator dense_to_sparse_coo, input argument must be a Tensor.")
-    if len(tensor.shape) != 2:
+    if len(_convert_shape(tensor.shape)) > 2:
         raise_value_error("Currently only support 2-D Tensor when converting to COOTensor.")
     indices = tensor.nonzero().astype("int32")
     values = gather_nd(tensor, indices)
-    return COOTensor(indices, values, tensor.shape)
+    return COOTensor(indices, values, _convert_shape(tensor.shape))
 
 
-def dense_to_sparse_csr(tensor):
+def dense_to_sparse_csr(tensor: Tensor) -> CSRTensor:
     """
     Convert a Tensor to CSRTensor.
 
@@ -395,11 +527,11 @@ def dense_to_sparse_csr(tensor):
     """
     if not isinstance(tensor, Tensor):
         raise_type_error("For functional operator dense_to_sparse_csr, input argument must be a Tensor.")
-    if len(tensor.shape) != 2:
+    if len(_convert_shape(tensor.shape)) > 2:
         raise_value_error("Currently only support 2-D Tensor when converting to CSRTensor.")
     indices = tensor.nonzero().astype("int32")
     _, _, indptr, indices, values = dense_to_csr(tensor, indices)
-    return CSRTensor(indptr, indices, values, tensor.shape)
+    return CSRTensor(indptr, indices, values, _convert_shape(tensor.shape))
 
 
 def make_sparse_tensor(indices, values, dense_shape):
@@ -409,11 +541,19 @@ def make_sparse_tensor(indices, values, dense_shape):
     return make_coo_tensor(indices, values, dense_shape)
 
 
+def make_row_tensor(indices, values, dense_shape):
+    """Call make_row_tensor_inner in this function."""
+    print_info("WARNING: 'RowTensor' is deprecated from version 2.0 and will be removed in a future version.")
+    return make_row_tensor_inner(indices, values, dense_shape)
+
+
 make_coo_tensor = Primitive('MakeCOOTensor')
 
 make_csr_tensor = Primitive('MakeCSRTensor')
 
-make_row_tensor = Primitive('MakeRowTensor')
+make_row_tensor_inner = Primitive('MakeRowTensor')
+
+make_map_parameter = Primitive('MakeMapParameter')
 
 row_tensor_get_values = Primitive('RowTensorGetValues')
 
@@ -426,7 +566,7 @@ row_tensor_add = Primitive('RowTensorAdd')
 
 @constexpr
 def _calc_out_shape(sp_input, concat_dim):
-    "calculating the COOTensor output shape in sparse_concat"
+    "calculating the COOTensor output shape in coo_concat"
     if isinstance(sp_input[0], tuple):
         out_shape_list = list(sp_input[0][2])
     else:
@@ -440,10 +580,10 @@ def _calc_out_shape(sp_input, concat_dim):
 
 
 @constexpr
-def _set_sparse_concat_input(sp_input):
+def _set_coo_concat_input(sp_input):
     "split COOTensor to normal tensor"
     if len(sp_input) < 2:
-        raise_value_error("For sparse_concat, not support COOTensor input number < 2.")
+        raise_value_error("For coo_concat, not support COOTensor input number < 2.")
     in_indices = []
     in_values = []
     in_shapes = []
@@ -459,21 +599,21 @@ def _set_sparse_concat_input(sp_input):
     return in_indices, in_values, in_shapes
 
 
-def sparse_concat(sp_input, concat_dim=0):
+def coo_concat(sp_input, concat_dim=0):
     """
     concatenates the input SparseTensor(COO format) along the specified dimension.
 
-    .. note:
-        demo API now, and only supported CPU
+    .. warning::
+        This is an experimental API that is subjected to change or deletion. Only supported on CPU now.
 
     Args:
-        sp_input (Union[list(COOTensor), tuple(COOTensor)) - the list of SparseTensor which need to concatenates.
+        sp_input (Union[list(COOTensor), tuple(COOTensor)]) - the list of SparseTensor which need to concatenates.
             for COOTensor input.
         concat_dim (scalar): decide the dimension to concatenation along.
             The value must be in range [-rank, rank), where rank is the number of dimensions in each input
             SparseTensor. Default is 0.
 
-    Outputs:
+    Returns:
         - **output** (COOtensor) - the result of concatenates the input SparseTensor along the
           specified dimension. OutShape: OutShape[non concat_dim] is equal to InShape[non concat_dim] and
           OutShape[concat_dim] is all input concat_dim axis shape accumulate.
@@ -486,39 +626,42 @@ def sparse_concat(sp_input, concat_dim=0):
         ``CPU``
 
     Examples:
-        >>> indices0 = Tensor([[0, 1], [1, 2]], dtype=mstype.int32)
+        >>> indices0 = Tensor([[0, 1], [1, 2]], dtype=mstype.int64)
         >>> values0 = Tensor([1, 2], dtype=mstype.int32)
         >>> shape0 = (3, 4)
         >>> input0 = COOTensor(indices0, values0, shape0)
-        >>> indices1 = Tensor([[0, 0], [1, 1]], dtype=mstype.int32)
+        >>> indices1 = Tensor([[0, 0], [1, 1]], dtype=mstype.int64)
         >>> values1 = Tensor([3, 4], dtype=mstype.int32)
         >>> shape1 = (3, 4)
         >>> input1 = COOTensor(indices1, values1, shape1)
         >>> concat_dim = 1
-        >>> out = F.sparse_concat((input0, input1), concat_dim)
+        >>> out = ops.coo_concat((input0, input1), concat_dim)
         >>> print(out)
-        COOTensor(shape=[3, 8], dtype=Int32, indices=Tensor(shape=[4, 2], dtype=Int32, value=
+        COOTensor(shape=[3, 8], dtype=Int32, indices=Tensor(shape=[4, 2], dtype=Int64, value=
         [[0 1]
          [0 4]
          [1 2]
          [1 5]]), values=Tensor(shape=[4], dtype=Int32, value=[1 3 2 4]))
     """
-    sparse_concat_op = SparseConcat(concat_dim)
-    in_indices, in_values, in_shapes = _set_sparse_concat_input(sp_input)
-    indices, values, _ = sparse_concat_op(in_indices, in_values, in_shapes)
+    coo_concat_op = SparseConcat(concat_dim)
+    in_indices, in_values, in_shapes = _set_coo_concat_input(sp_input)
+    indices, values, _ = coo_concat_op(in_indices, in_values, in_shapes)
     out_shape = _calc_out_shape(sp_input, concat_dim)
     return COOTensor(indices, values, out_shape)
 
 
-def sparse_add(x1, x2, thresh):
+def coo_add(x1: COOTensor, x2: COOTensor, thresh: Tensor) -> COOTensor:
     """
-    Computes the sum of x1(COOTensor) and x2(COOTensor).
+    Computes the sum of x1(COOTensor) and x2(COOTensor), and return a new COOTensor
+    based on the computed result and `thresh`.
 
     Args:
         x1 (COOTensor): the first COOTensor to sum.
         x2 (COOTensor): the second COOTensor to sum.
         thresh (Tensor): A 0-D Tensor, represents the magnitude threshold that determines
-            if an output value/index pair pair take space.
+            if an output value/index pair take place. Its dtype
+            should match that of the values if they are real. If output's
+            value is less than the `thresh`, it will vanish.
 
     Returns:
         A COOTensor, the result of sum.
@@ -542,7 +685,7 @@ def sparse_add(x1, x2, thresh):
         TypeError: If (x1/x2)'s value's type is not matched with thresh's type.
 
     Supported Platforms:
-        ``CPU`` ``GPU``
+        ``GPU`` ``CPU``
 
     Examples:
         >>> from mindspore import Tensor, COOTensor
@@ -558,12 +701,16 @@ def sparse_add(x1, x2, thresh):
         >>> shape1 = (3, 4)
         >>> input1 = COOTensor(indics1, values1, shape1)
         >>> thres = Tensor(0, dtype=mstype.int32)
-        >>> out = ops.sparse_add(input0, input1, thres)
+        >>> out = ops.coo_add(input0, input1, thres)
         >>> print(out)
-        COOTensor(shape = [3, 4], dtype = Int32, indices=Tensor(shape=[4, 2],
-        dtype = Int64, value=[[0 0], [0 1], [1 1], [1 2]]),  values=Tensor(shape[4],
-        dtype=Int32, value=[3 1 4 2]))
+        COOTensor(shape=[3, 4], dtype=Int32, indices=Tensor(shape=[4, 2], dtype=Int64, value=
+        [[0 0]
+         [0 1]
+         [1 1]
+         [1 2]]), values=Tensor(shape=[4], dtype=Int32, value=[3 1 4 2]))
     """
+    x1, x2 = promote_coo(x1, x2)
+    thresh = thresh.astype(x1.dtype)
     x1_indices = x1.indices
     x1_values = x1.values
     x2_indices = x2.indices
@@ -574,36 +721,36 @@ def sparse_add(x1, x2, thresh):
     return COOTensor(indices, values, x1.shape)
 
 
-def csr_softmax(logits, dtype):
+def csr_softmax(logits: CSRTensor, dtype: mstype):
     """
     Calculates the softmax of a CSRTensorMatrix.
 
     Args:
-        logits (CSRTensor): Sparse CSR Tensor.
-        dtype (dtype): Data type.
+        logits (CSRTensor): Input sparse CSRTensor.
+        dtype (dtype): Input data type.
 
     Returns:
-        CSRTensor, a csr_tensor containing
+        CSRTensor, a CSRTensor containing
 
-        - **indptr** -  indicates the start and end point for `values` in each row.
-        - **indices** - the column positions of all non-zero values of the input.
-        - **values** - the non-zero values of the dense tensor.
-        - **shape** - the shape of the csr_tensor.
+        - **indptr** - Indicates the start and end point for non-zero values in each row.
+        - **indices** - The column positions of all non-zero values of the input.
+        - **values** - The non-zero values of the dense tensor.
+        - **shape** - The shape of the CSRTensor.
 
     Supported Platforms:
         ``GPU`` ``CPU``
 
     Examples:
         >>> import mindspore as ms
+        >>> import mindspore.ops as ops
         >>> import mindspore.common.dtype as mstype
         >>> from mindspore import Tensor, CSRTensor
-        >>> from mindspore.ops.function import csr_softmax
         >>> logits_indptr = Tensor([0, 4, 6], dtype=mstype.int32)
         >>> logits_indices = Tensor([0, 2, 3, 4, 3, 4], dtype=mstype.int32)
         >>> logits_values = Tensor([1, 2, 3, 4, 1, 2], dtype=mstype.float32)
         >>> shape = (2, 6)
         >>> logits = CSRTensor(logits_indptr, logits_indices, logits_values, shape)
-        >>> out = csr_softmax(logits, dtype=mstype.float32)
+        >>> out = ops.csr_softmax(logits, dtype=mstype.float32)
         >>> print(out)
         CSRTensor(shape=[2, 6], dtype=Float32, indptr=Tensor(shape=[3], dtype=Int32, value=[0 4 6]),
                        indices=Tensor(shape=[6], dtype=Int32, value=[0 2 3 4 3 4]),
@@ -621,23 +768,35 @@ def csr_softmax(logits, dtype):
     return CSRTensor(indptr=indptr, indices=indices, values=values, shape=output_shape)
 
 
-def csr_add(a, b, alpha, beta):
+def csr_add(a: CSRTensor, b: CSRTensor, alpha: Tensor, beta: Tensor) -> CSRTensor:
     """
-    Returns alpha * csr_a + beta * csr_b where both csr_a and csr_b are CSRTensor, alpha and beta are both Tensor.
+    Computes the linear combination of two input CSRTensors a and b.
+
+    .. math::
+
+        out = alpha * a + beta * b
+
+    where both :math:`a` and :math:`b` are CSRTensor, :math:`alpha` and :math:`beta` are both Tensor
+
+    Note:
+        The user need to ensure that the input sparse matrix is legal.
+        Otherwise, the behavior of the operator is undefined.
+        For example, when there are multiple elements in the same position, the
+        operator may return an error of fail execute.
 
     Args:
-        a (CSRTensor): Sparse CSR Tensor.
-        b (CSRTensor): Sparse CSR Tensor.
+        a (CSRTensor): Input sparse CSRTensor.
+        b (CSRTensor): Input sparse CSRTensor.
         alpha(Tensor): Dense Tensor, its shape must be able to broadcast to a.
         beta(Tensor): Dense Tensor, its shape must be able to broadcast to b.
 
     Returns:
-        CSRTensor, a csr_tensor containing the following parts.
+        CSRTensor, a CSRTensor containing the following parts.
 
-        - **indptr** -  indicates the start and end point for `values` in each row.
-        - **indices** - the column positions of all non-zero values of the input.
-        - **values** - the non-zero values of the dense tensor.
-        - **shape** - the shape of the csr_tensor.
+        - **indptr** -  Indicates the start and end point for non-zero values in each row.
+        - **indices** - The column positions of all non-zero values of the input.
+        - **values** - The non-zero values of the dense tensor.
+        - **shape** - The shape of the CSRTensor.
 
     Supported Platforms:
         ``GPU`` ``CPU``
@@ -645,7 +804,7 @@ def csr_add(a, b, alpha, beta):
     Examples:
         >>> import mindspore.common.dtype as mstype
         >>> from mindspore import Tensor, CSRTensor
-        >>> from mindspore.ops.functional import csr_add
+        >>> import mindspore.ops as ops
         >>> a_indptr = Tensor([0, 1, 2], dtype=mstype.int32)
         >>> a_indices = Tensor([0, 1], dtype=mstype.int32)
         >>> a_values = Tensor([1, 2], dtype=mstype.float32)
@@ -657,7 +816,7 @@ def csr_add(a, b, alpha, beta):
         >>> beta = Tensor(1, mstype.float32)
         >>> csra = CSRTensor(a_indptr, a_indices, a_values, shape)
         >>> csrb = CSRTensor(b_indptr, b_indices, b_values, shape)
-        >>> out = csr_add(csra, csrb, alpha, beta)
+        >>> out = ops.csr_add(csra, csrb, alpha, beta)
         >>> print(out)
         CSRTensor(shape=[2,6], dtype=Float32,
                   indptr=Tensor(shape=[3], dtype=Int32, value = [0, 1, 2]),
@@ -669,10 +828,13 @@ def csr_add(a, b, alpha, beta):
     if not isinstance(alpha, Tensor) or not isinstance(beta, Tensor):
         raise_type_error("For functional operator csr_add, both inputs alpha and beta must be Tensor.")
     csr_add_op = SparseMatrixAdd()
-    a_batch_pointers = make_tensor([0, a.values.shape[0]], mstype.int32)
-    b_batch_pointers = make_tensor([0, b.values.shape[0]], mstype.int32)
-    a_shape = make_tensor(a.shape, mstype.int32)
-    b_shape = make_tensor(b.shape, mstype.int32)
+    a, b = promote_csr(a, b)
+    alpha = alpha.astype(a.dtype)
+    beta = beta.astype(a.dtype)
+    a_batch_pointers = ops.concat((make_tensor([0]), ops.TensorShape()(a.values))).astype(a.indptr.dtype)
+    b_batch_pointers = ops.concat((make_tensor([0]), ops.TensorShape()(b.values))).astype(b.indptr.dtype)
+    a_shape = make_tensor(a.shape, a.indptr.dtype)
+    b_shape = make_tensor(b.shape, b.indptr.dtype)
     _, _, indptr, indices, values = csr_add_op(a_shape, a_batch_pointers, a.indptr, a.indices, a.values,
                                                b_shape, b_batch_pointers, b.indptr, b.indices, b.values,
                                                alpha, beta)
@@ -687,6 +849,7 @@ __all__ = [
     'coo_tensor_get_values',
     'csr_div',
     'csr_gather',
+    'csr_mm',
     'csr_mul',
     'csr_mv',
     'csr_reduce_sum',
@@ -702,12 +865,14 @@ __all__ = [
     'make_coo_tensor',
     'make_csr_tensor',
     'make_row_tensor',
+    'make_row_tensor_inner',
+    'make_map_parameter',
     'row_tensor_get_values',
     'row_tensor_get_indices',
     'row_tensor_get_dense_shape',
     'row_tensor_add',
-    'sparse_add',
-    'sparse_concat',
+    'coo_add',
+    'coo_concat',
     'csr_add',
     'csr_softmax',
     'csr_to_dense'