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