• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# coding: utf-8
2
3# Copyright 2020-2021 Huawei Technologies Co., Ltd
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15
16# limitations under the License.
17# ============================================================================
18
19"""Operators for sparse operators."""
20
21from ..._checkparam import Validator as validator
22from ...common import dtype as mstype
23from ..primitive import PrimitiveWithInfer, prim_attr_register
24
25
26class SparseToDense(PrimitiveWithInfer):
27    """
28    Converts a sparse representation into a dense tensor.
29
30    Inputs:
31        - **indices** (Tensor) - A 2-D Tensor, represents the position of the element in the sparse tensor.
32          Support int32, int64, each element value should be a non-negative int number. The shape is :math:`(n, 2)`.
33        - **values** (Tensor) - A 1-D Tensor, represents the value corresponding to the position in the `indices`.
34          The shape should be :math:`(n,)`.
35        - **sparse_shape** (tuple(int)) - A positive int tuple which specifies the shape of sparse tensor,
36          should have 2 elements, represent sparse tensor shape is :math:`(N, C)`.
37
38    Returns:
39        Tensor, converted from sparse tensor. The dtype is same as `values`, and the shape is `sparse_shape`.
40
41    Raises:
42        TypeError: If the dtype of `indices` is neither int32 nor int64.
43        ValueError: If `sparse_shape`, shape of `indices and shape of `values` don't meet the parameter description.
44
45    Supported Platforms:
46        ``CPU``
47
48    Examples:
49        >>> indices = Tensor([[0, 1], [1, 2]])
50        >>> values = Tensor([1, 2], dtype=ms.float32)
51        >>> sparse_shape = (3, 4)
52        >>> sparse_to_dense = ops.SparseToDense()
53        >>> out = sparse_to_dense(indices, values, sparse_shape)
54        >>> print(out)
55        [[0 1 0 0]
56         [0 0 2 0]
57         [0 0 0 0]]
58    """
59
60    @prim_attr_register
61    def __init__(self):
62        """Initialize SparseToDense."""
63        self.init_prim_io_names(inputs=['indices', 'values', 'dense_shape'], outputs=['output'])
64
65    def __infer__(self, indices, values, sparse_shape):
66        validator.check_tensor_dtype_valid('indices', indices['dtype'], [mstype.int32, mstype.int64], self.name)
67        validator.check_tensor_dtype_valid('values', values['dtype'], mstype.number_type + (mstype.bool_,), self.name)
68        indices_shape = indices['shape']
69        if len(indices_shape) != 2:
70            raise ValueError(f"For '{self.name}', the 'indices' must be a 2-D tensor, "
71                             f"but got 'indices' shape: {indices_shape}.")
72        values_shape = values['shape']
73        if len(values_shape) != 1 or values_shape[0] != indices_shape[0]:
74            raise ValueError(f"For '{self.name}', the 'values' must be a 1-D tensor and the first dimension length "
75                             f"must be equal to the first dimension length of 'indices', "
76                             f"but got 'indices' shape: {indices_shape}, 'values' shape: {values_shape}.")
77        sparse_shape_v = sparse_shape['value']
78        for i in sparse_shape_v:
79            if isinstance(i, bool) or not isinstance(i, int) or i <= 0:
80                raise ValueError(f"For '{self.name}', all elements in 'sparse_shape' must be "
81                                 f"positive int number, but got 'sparse_shape': {sparse_shape_v}.")
82        if len(sparse_shape_v) != indices_shape[1]:
83            raise ValueError(f"For '{self.name}', the length of 'sparse_shape' should be equal to the second dimension "
84                             f"length of 'indices', but got the second dimension length of 'indices': "
85                             f"{indices_shape[1]}, length of 'sparse_shape': {len(sparse_shape_v)}.")
86        out = {'shape': sparse_shape['value'],
87               'dtype': values['dtype'],
88               'value': None}
89        return out
90
91
92class SparseTensorDenseMatmul(PrimitiveWithInfer):
93    """
94    Multiplies sparse matrix `A` by dense matrix `B`.
95    The rank of sparse matrix and dense matrix must be equal to `2`.
96
97    Args:
98        adjoint_st (bool): If true, sparse tensor is transposed before multiplication. Default: False.
99        adjoint_dt (bool): If true, dense tensor is transposed before multiplication. Default: False.
100
101    Inputs:
102        - **indices** (Tensor) - A 2-D Tensor, represents the position of the element in the sparse tensor.
103          Support int32, int64, each element value should be a non-negative int number. The shape is :math:`(n, 2)`.
104        - **values** (Tensor) - A 1-D Tensor, represents the value corresponding to the position in the `indices`.
105          Support float16, float32, float64, int32, int64. The shape should be :math:`(n,)`.
106        - **sparse_shape** (tuple(int)) - A positive int tuple which specifies the shape of sparse tensor,
107          should have 2 elements, represent sparse tensor shape is :math:`(N, C)`.
108        - **dense** (Tensor) - A 2-D Tensor, the dtype is same as `values`.
109          If `adjoint_st` is False and `adjoint_dt` is False, the shape must be :math:`(C, M)`.
110          If `adjoint_st` is False and `adjoint_dt` is True, the shape must be :math:`(M, C)`.
111          If `adjoint_st` is True and `adjoint_dt` is False, the shape must be :math:`(N, M)`.
112          If `adjoint_st` is True and `adjoint_dt` is True, the shape must be :math:`(M, N)`.
113
114    Outputs:
115        Tensor, the dtype is the same as `values`.
116        If `adjoint_st` is False, the shape is :math:`(N, M)`.
117        If `adjoint_st` is True, the shape is :math:`(C, M)`.
118
119    Raises:
120        TypeError: If the type of `adjoint_st` or `adjoint_dt` is not bool, or the dtype of `indices`,
121            dtype of `values` and dtype of `dense` don't meet the parameter description.
122        ValueError: If `sparse_shape`, shape of `indices, shape of `values`,
123            and shape of `dense` don't meet the parameter description.
124
125    Supported Platforms:
126        ``CPU``
127
128    Examples:
129        >>> indices = Tensor([[0, 1], [1, 2]], dtype=ms.int32)
130        >>> values = Tensor([1, 2], dtype=ms.float32)
131        >>> sparse_shape = (3, 4)
132        >>> dense = Tensor([[1,1], [2,2], [3,3 ], [4, 4]], dtype=ms.float32)
133        >>> sparse_dense_matmul = ops.SparseTensorDenseMatmul()
134        >>> out = sparse_dense_matmul(indices, values, sparse_shape, dense)
135        >>> print(out)
136        [[2 2]
137         [6 6]
138         [0 0]]
139    """
140
141    @prim_attr_register
142    def __init__(self, adjoint_st=False, adjoint_dt=False):
143        """Initialize SparseTensorDenseMatmul"""
144        self.adjoint_st = adjoint_st
145        self.adjoint_dt = adjoint_dt
146        self.init_prim_io_names(inputs=['indices', 'values', 'sparse_shape', 'dense'],
147                                outputs=['output'])
148        self.add_prim_attr('adjoint_st', self.adjoint_st)
149        self.add_prim_attr('adjoint_dt', self.adjoint_dt)
150        validator.check_value_type("adjoint_st", adjoint_st, [bool], self.name)
151        validator.check_value_type("adjoint_dt", adjoint_dt, [bool], self.name)
152
153    def __infer__(self, indices, values, sparse_shape, dense):
154        validator.check_tensor_dtype_valid('indices', indices['dtype'], [mstype.int32, mstype.int64], self.name)
155        valid_types = (mstype.float16, mstype.float32, mstype.float64, mstype.int32, mstype.int64)
156        args = {'values': values['dtype'], 'dense': dense['dtype']}
157        validator.check_tensors_dtypes_same_and_valid(args, valid_types, self.name)
158        indices_shape = indices['shape']
159        if len(indices_shape) != 2 or indices_shape[1] != 2:
160            raise ValueError(f"For '{self.name}', the 'indices' must be a 2-D tensor and "
161                             f"the second dimension length must be 2, but got 'indices' shape: {indices_shape}.")
162        values_shape = values['shape']
163        if len(values_shape) != 1 or values_shape[0] != indices_shape[0]:
164            raise ValueError(f"For '{self.name}', the 'values' must be a 1-D tensor and "
165                             f"the first dimension length must be equal to the first dimension length of 'indices', "
166                             f"but got 'indices' shape: {indices_shape}, 'values' shape: {values_shape}.")
167        a_shape = sparse_shape['value'][::-1] if self.adjoint_st else sparse_shape['value']
168        b_shape = dense['shape'][::-1] if self.adjoint_dt else dense['shape']
169        for i in a_shape:
170            if isinstance(i, bool) or not isinstance(i, int) or i <= 0:
171                raise ValueError(f"For '{self.name}', all elements in 'sparse_shape' must be "
172                                 f"positive int number, but got 'sparse_shape': {a_shape}.")
173        if len(a_shape) != 2 or len(b_shape) != 2:
174            raise ValueError(f"For '{self.name}', both the length of 'sparse_shape' and the tensor "
175                             f"rank of 'dense' should be equal to 2, but got the length of "
176                             f"'sparse_shape': {len(a_shape)}, "
177                             f"the tensor rank of 'dense': {len(b_shape)}.")
178        if a_shape[1] != b_shape[0]:
179            raise ValueError(f"For '{self.name}', the second dimension length of 'sparse_shape' must be equal to the "
180                             f"first dimension length of 'dense', but got "
181                             f"the tensor shape of 'sparse': {a_shape} and the tensor shape of 'dense': {b_shape}. "
182                             f"Don't meet the condition for matmul")
183        out_shape = [a_shape[0], b_shape[1]]
184        out = {'shape': tuple(out_shape),
185               'dtype': values['dtype'],
186               'value': None}
187        return out
188