mindspore 2.0.0a0__cp37-none-any.whl → 2.0.0rc1__cp37-none-any.whl

mindspore/ops/function/sparse_func.py

@@ -32,6 +32,7 @@ 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.ops.composite.multitype_ops._constexpr_utils import raise_value_error, raise_type_error, make_tensor,\
     promote_binary_dtype
@@ -57,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:
@@ -65,11 +80,13 @@ def is_scalar(tensor):
 
 
 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)
@@ -79,6 +96,7 @@ def promote_csr(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)
@@ -118,7 +136,7 @@ def coalesce(x_indices: Tensor, x_values: Tensor, x_shape: Tensor) -> Tuple[Tens
         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
@@ -175,28 +193,23 @@ def csr_div(x: CSRTensor, y: Tensor) -> Tensor:
             raise_value_error("dense tensor cannot broadcast to the sparse tensor.")
         return (x.values / y).reshape(x.values.shape)
     x_values, y = promote_tensor(x.values, y)
-    return _csr_ops.CSRDiv()(x.indptr, x.indices, x_values, x.shape, 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: CSRTensor, y: Tensor) -> Tensor:
+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``
@@ -208,7 +221,8 @@ def csr_mul(x: CSRTensor, y: Tensor) -> Tensor:
             raise_value_error("dense tensor cannot broadcast to the sparse tensor.")
         return (x.values * y).reshape(x.values.shape)
     x_values, y = promote_tensor(x.values, y)
-    return _csr_ops.CSRMul()(x.indptr, x.indices, x_values, x.shape, 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: CSRTensor, dense: Tensor) -> Tensor:
@@ -242,10 +256,10 @@ def csr_mm(a: CSRTensor, b: CSRTensor, trans_a: bool = False, trans_b: bool = Fa
     Args:
         a (CSRTensor): Sparse CSR Tensor, rank should be 2.
         b (CSRTensor): Sparse CSR Tensor, rank should be 2.
-        trans_a (bool): whether to transpose CSRTensor a. Default: False.
-        trans_b (bool): whether to transpose CSRTensor b. Default: False.
-        adjoint_a (bool): whether to adjoint CSRTensor a. Default: False.
-        adjoint_b (bool): whether to adjoint CSRTensor b. Default: False.
+        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.
@@ -277,35 +291,32 @@ def csr_mm(a: CSRTensor, b: CSRTensor, trans_a: bool = False, trans_b: bool = Fa
         >>> print(c.indices)
         [0 0]
     """
-    if isinstance(a, CSRTensor) and isinstance(b, CSRTensor):
-        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)
-        m, a2 = a.shape
-        b1, k = b.shape
-        if trans_a or adjoint_a:
-            m = a2
-        if trans_b or adjoint_b:
-            k = b1
-        return CSRTensor(c_indptr, c_indices, c_values, (m, k))
-    raise_type_error("For functional operator csr_mm, inputs a and b must be type of CSRTensor currently.")
-    return None
+    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:
@@ -359,12 +370,12 @@ def csr_to_coo(tensor: CSRTensor) -> COOTensor:
     """
     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: CSRTensor) -> Tensor:
@@ -399,12 +410,12 @@ def csr_to_dense(csr_tensor: CSRTensor) -> 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
+    shape = _convert_shape(csr_tensor.shape)
     dense_shape = Tensor(shape, dtype=mstype.int32)
-    batch_pointers = Tensor([0, -1], 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
@@ -470,11 +481,11 @@ def dense_to_sparse_coo(tensor: Tensor) -> COOTensor:
     """
     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: Tensor) -> CSRTensor:
@@ -516,11 +527,11 @@ def dense_to_sparse_csr(tensor: Tensor) -> CSRTensor:
     """
     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):
@@ -555,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:
@@ -569,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 = []
@@ -588,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.
@@ -615,31 +626,31 @@ 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: COOTensor, x2: COOTensor, thresh: Tensor) -> COOTensor:
+def coo_add(x1: COOTensor, x2: COOTensor, thresh: Tensor) -> COOTensor:
     """
     Computes the sum of x1(COOTensor) and x2(COOTensor), and return a new COOTensor
     based on the computed result and `thresh`.
@@ -648,7 +659,7 @@ def sparse_add(x1: COOTensor, x2: COOTensor, thresh: Tensor) -> COOTensor:
         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 take space. Its dtype
+            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.
 
@@ -674,7 +685,7 @@ def sparse_add(x1: COOTensor, x2: COOTensor, thresh: Tensor) -> COOTensor:
         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
@@ -690,11 +701,13 @@ def sparse_add(x1: COOTensor, x2: COOTensor, thresh: Tensor) -> COOTensor:
         >>> 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)
@@ -757,7 +770,13 @@ def csr_softmax(logits: CSRTensor, dtype: mstype):
 
 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.
@@ -830,6 +849,7 @@ __all__ = [
     'coo_tensor_get_values',
     'csr_div',
     'csr_gather',
+    'csr_mm',
     'csr_mul',
     'csr_mv',
     'csr_reduce_sum',
@@ -851,8 +871,8 @@ __all__ = [
     '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'