1# Copyright 2022 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15import pytest 16import mindspore.context as context 17import mindspore.nn as nn 18import mindspore as ms 19import mindspore.ops.operations.sparse_ops as P 20from mindspore import Tensor 21 22 23class Net(nn.Cell): 24 25 def __init__(self): 26 super(Net, self).__init__() 27 self.op = P.SparseMatrixAdd() 28 29 def construct(self, a_shape, a_batch_pointer, a_indptr, a_indices, 30 a_values, b_shape, b_batch_pointer, b_indptr, b_indices, 31 b_values, alpha, beta): 32 return self.op(a_shape, a_batch_pointer, a_indptr, a_indices, a_values, 33 b_shape, b_batch_pointer, b_indptr, b_indices, b_values, 34 alpha, beta) 35 36 37def dyn_case(): 38 net = Net() 39 40 a_indptr_dyn = Tensor(shape=[None], dtype=ms.int32) 41 a_indices_dyn = Tensor(shape=[None], dtype=ms.int32) 42 a_values_dyn = Tensor(shape=[None], dtype=ms.float32) 43 a_pointers_dyn = Tensor(shape=[None], dtype=ms.int32) 44 shape_dyn = Tensor(shape=[None], dtype=ms.int32) 45 b_indptr_dyn = Tensor(shape=[None], dtype=ms.int32) 46 b_indices_dyn = Tensor(shape=[None], dtype=ms.int32) 47 b_values_dyn = Tensor(shape=[None], dtype=ms.float32) 48 b_pointers_dyn = Tensor(shape=[None], dtype=ms.int32) 49 alpha = Tensor(1, ms.float32) 50 beta = Tensor(1, ms.float32) 51 net.set_inputs(shape_dyn, a_pointers_dyn, a_indptr_dyn, a_indices_dyn, 52 a_values_dyn, shape_dyn, b_pointers_dyn, b_indptr_dyn, 53 b_indices_dyn, b_values_dyn, alpha, beta) 54 55 a_indptr = Tensor([0, 1, 2], dtype=ms.int32) 56 a_indices = Tensor([0, 1], dtype=ms.int32) 57 a_values = Tensor([1, 2], dtype=ms.float32) 58 a_pointers = Tensor([0, a_values.shape[0]], dtype=ms.int32) 59 shape = Tensor([2, 6], dtype=ms.int32) 60 b_indptr = Tensor([0, 1, 2], dtype=ms.int32) 61 b_indices = Tensor([0, 1], dtype=ms.int32) 62 b_values = Tensor([1, 2], dtype=ms.float32) 63 b_pointers = Tensor([0, b_values.shape[0]], dtype=ms.int32) 64 out = net(shape, a_pointers, a_indptr, a_indices, a_values, shape, 65 b_pointers, b_indptr, b_indices, b_values, alpha, beta) 66 67 exepct_shapes = [(2,), (2,), (3,), (2,), (2,)] 68 for i in range(5): 69 assert out[i].asnumpy().shape == exepct_shapes[i] 70 71 72@pytest.mark.level0 73@pytest.mark.platform_x86_gpu 74@pytest.mark.env_onecard 75def test_sparse_matrix_add_dyn(): 76 """ 77 Feature: test SparseMatrixAdd in gpu. 78 Description: test the ops in dynamic case. 79 Expectation: success. 80 """ 81 context.set_context(mode=context.GRAPH_MODE, device_target='GPU') 82 dyn_case() 83 context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU') 84 dyn_case() 85