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# ============================================================================ 15"""smoke tests for COO operations""" 16 17import pytest 18import numpy as np 19 20from mindspore import Tensor, COOTensor, jit, nn, ops 21from mindspore.common import dtype as mstype 22from mindspore.ops import functional as F 23 24from .sparse_utils import get_platform, compare_res, compare_coo 25 26 27@pytest.mark.level1 28@pytest.mark.platform_arm_ascend_training 29@pytest.mark.platform_x86_ascend_training 30@pytest.mark.platform_x86_gpu_training 31@pytest.mark.platform_x86_cpu 32@pytest.mark.env_onecard 33def test_make_coo(): 34 """ 35 Feature: Test COOTensor Constructor in Graph and PyNative. 36 Description: Test COOTensor(indices, values, shape) and COOTensor(COOTensor) 37 Expectation: Success. 38 """ 39 indices = Tensor([[0, 1], [1, 2]]) 40 values = Tensor([1, 2], dtype=mstype.float32) 41 dense_shape = (3, 4) 42 43 def test_pynative(): 44 return COOTensor(indices, values, dense_shape) 45 test_graph = jit(test_pynative) 46 47 coo1 = test_pynative() 48 coo2 = test_graph() 49 compare_coo(coo1, coo2) 50 coo3 = COOTensor(coo_tensor=coo2) 51 compare_coo(coo3, coo2) 52 53 54@pytest.mark.level0 55@pytest.mark.platform_arm_ascend_training 56@pytest.mark.platform_x86_ascend_training 57@pytest.mark.platform_x86_gpu_training 58@pytest.mark.env_onecard 59def test_coo_tensor_with_control_if(): 60 """ 61 Feature: Test COOTensor in if. 62 Description: Test COOTensor computation in while loop. 63 Expectation: Success. 64 """ 65 class COOTensorValuesDouble(nn.Cell): 66 67 def construct(self, x): 68 indices = x.indices 69 values = x.values * 2 70 shape = x.shape 71 return COOTensor(indices, values, shape) 72 73 class COOTensorValuesAdd2(nn.Cell): 74 75 def construct(self, x): 76 indices = x.indices 77 values = x.values + 2 78 shape = x.shape 79 return COOTensor(indices, values, shape) 80 81 class COOTensorWithControlIf(nn.Cell): 82 def __init__(self, shape): 83 super(COOTensorWithControlIf, self).__init__() 84 self.op1 = COOTensorValuesDouble() 85 self.op2 = COOTensorValuesAdd2() 86 self.shape = shape 87 88 def construct(self, a, b, indices, values): 89 x = COOTensor(indices, values, self.shape) 90 if a > b: 91 x = self.op1(x) 92 else: 93 x = self.op2(x) 94 return x.indices, x.values, x.shape 95 96 a = Tensor(0, mstype.int32) 97 b = Tensor(2, mstype.int32) 98 indices = Tensor([[0, 1], [1, 2]]) 99 values = Tensor([1, 2], dtype=mstype.float32) 100 shape = (3, 4) 101 net = COOTensorWithControlIf(shape) 102 out = net(a, b, indices, values) 103 assert np.allclose(out[0].asnumpy(), indices.asnumpy(), .0, .0) 104 assert np.allclose(out[1].asnumpy(), values.asnumpy() + 2, .0, .0) 105 assert out[2] == shape 106 107 108@pytest.mark.level2 109@pytest.mark.platform_arm_ascend_training 110@pytest.mark.platform_x86_ascend_training 111@pytest.mark.platform_x86_gpu_training 112@pytest.mark.env_onecard 113def test_coo_tensor_in_while(): 114 """ 115 Feature: Test COOTensor in while loop. 116 Description: Test COOTensor computation in while loop. 117 Expectation: Success. 118 """ 119 class COOTensorWithControlWhile(nn.Cell): 120 def __init__(self, shape): 121 super(COOTensorWithControlWhile, self).__init__() 122 self.shape = shape 123 124 @jit 125 def construct(self, a, b, indices, values): 126 x = COOTensor(indices, values, self.shape) 127 while a > b: 128 x = COOTensor(indices, values, self.shape) 129 b = b + 1 130 return x 131 a = Tensor(3, mstype.int32) 132 b = Tensor(0, mstype.int32) 133 indices = Tensor([[0, 1], [1, 2]]) 134 values = Tensor([1, 2], dtype=mstype.float32) 135 shape = (3, 4) 136 net = COOTensorWithControlWhile(shape) 137 out = net(a, b, indices, values) 138 assert np.allclose(out.indices.asnumpy(), indices.asnumpy(), .0, .0) 139 assert np.allclose(out.values.asnumpy(), values.asnumpy(), .0, .0) 140 assert out.shape == shape 141 142 143@pytest.mark.level1 144@pytest.mark.platform_x86_gpu_training 145@pytest.mark.platform_x86_cpu 146@pytest.mark.env_onecard 147def test_coo_method(): 148 """ 149 Feature: Test coo tensor methods. 150 Description: Test coo_tensor.to_csr(), coo_tensor.to_dense(). 151 Expectation: Success. 152 """ 153 if get_platform() != "linux": 154 return 155 156 class COOToCSRNet(nn.Cell): 157 def construct(self, coo_tensor): 158 return coo_tensor.to_csr() 159 160 class COOToDenseNet(nn.Cell): 161 def construct(self, coo_tensor): 162 return coo_tensor.to_dense() 163 164 indices = Tensor([[1, 2], [0, 1]], dtype=mstype.int32) 165 values = Tensor([2, 1], dtype=mstype.float32) 166 shape = (3, 4) 167 coo_tensor = COOTensor(indices, values, shape) 168 169 to_csr_output = COOToCSRNet()(coo_tensor) 170 to_csr_expect_1 = np.array([0, 1, 2, 2], dtype=np.int32) 171 to_csr_expect_2 = np.array([1, 2], dtype=np.int32) 172 to_csr_expect_3 = np.array([1, 2], dtype=np.float32) 173 assert np.allclose(to_csr_output.indptr.asnumpy(), to_csr_expect_1) 174 assert np.allclose(to_csr_output.indices.asnumpy(), to_csr_expect_2) 175 assert np.allclose(to_csr_output.values.asnumpy(), to_csr_expect_3) 176 177 to_dense_output = COOToDenseNet()(coo_tensor) 178 to_dense_expect = np.array( 179 [[0., 1., 0., 0.], [0., 0., 2., 0.], [0., 0., 0., 0.]], dtype=np.float32) 180 assert np.allclose(to_dense_output.asnumpy(), to_dense_expect) 181 182 183@pytest.mark.level1 184@pytest.mark.platform_x86_gpu_training 185@pytest.mark.env_onecard 186def test_coo_coalesce(): 187 """ 188 Feature: Test coo tensor coalesce methods. 189 Description: Test coo_tensor.coalesce() 190 Expectation: Success. 191 """ 192 if get_platform() != "linux": 193 return 194 195 class COOCoalesce(nn.Cell): 196 def construct(self, coo_tensor): 197 return coo_tensor.coalesce() 198 199 indices = Tensor([[0, 0, 1], [1, 1, 2]], dtype=mstype.int64) 200 values = Tensor([1, 5, 4], dtype=mstype.float32) 201 shape = (3, 3) 202 coo_tensor = COOTensor(indices.transpose(), values, shape) 203 204 coalesce_output = COOCoalesce()(coo_tensor) 205 res_coo = coo_tensor.coalesce() 206 207 expect_indices = np.array([[0, 1], [1, 2]], dtype=np.int64) 208 expect_values = np.array([6., 4.], dtype=np.float32) 209 assert np.allclose(expect_indices, res_coo.indices.asnumpy()) 210 assert np.allclose(expect_values, res_coo.values.asnumpy()) 211 assert np.allclose(expect_indices, coalesce_output.indices.asnumpy()) 212 assert np.allclose(expect_values, coalesce_output.values.asnumpy()) 213 214 215@pytest.mark.level2 216@pytest.mark.platform_x86_gpu_training 217@pytest.mark.platform_x86_cpu 218@pytest.mark.env_onecard 219def test_dtype_coo_tensor(): 220 """ 221 Feature: Test F.dtype with COOTensor. 222 Description: Test: F.dtype(x), x.dtype. 223 Expectation: Success. 224 """ 225 if get_platform() != "linux": 226 return 227 indices = Tensor([[0, 1], [1, 2]]) 228 values = Tensor([1, 2], dtype=mstype.float32) 229 shape = (3, 4) 230 231 def pynative_test(): 232 x = COOTensor(indices, values, shape) 233 return F.dtype(x), x.dtype 234 graph_test = jit(pynative_test) 235 236 out1, out2 = pynative_test() 237 out3, out4 = graph_test() 238 assert out1 in [mstype.float32] 239 assert out2 in [mstype.float32] 240 assert out3 in [mstype.float32] 241 assert out4 in [mstype.float32] 242 243 244@pytest.mark.level1 245@pytest.mark.platform_arm_ascend_training 246@pytest.mark.platform_x86_ascend_training 247@pytest.mark.platform_x86_gpu_training 248@pytest.mark.platform_x86_cpu 249@pytest.mark.env_onecard 250def test_coo_attr(): 251 """ 252 Feature: Test COOTensor GetAttr in Graph and PyNative. 253 Description: Test COOTensor.indices, COOTensor.values, COOTensor.shape. 254 Expectation: Success. 255 """ 256 if get_platform() != "linux": 257 return 258 indices = Tensor([[0, 1], [1, 2]]) 259 values = Tensor([1, 2], dtype=mstype.float32) 260 shape = (3, 4) 261 coo = COOTensor(indices, values, shape) 262 263 def test_pynative_1(): 264 return coo.indices, coo.values, coo.shape 265 266 def test_pynative_2(): 267 return coo.astype(mstype.int32) 268 269 def test_pynative_3(): 270 return coo.to_tuple() 271 272 test_graph_1 = jit(test_pynative_1) 273 test_graph_2 = jit(test_pynative_2) 274 test_graph_3 = jit(test_pynative_3) 275 276 py_indices, py_values, py_shape = test_pynative_1() 277 py_coo = test_pynative_2() 278 py_tuple = test_pynative_3() 279 280 g_indices, g_values, g_shape = test_graph_1() 281 g_coo = test_graph_2() 282 g_tuple = test_graph_3() 283 284 coo1 = COOTensor(py_indices, py_values, py_shape) 285 coo2 = COOTensor(g_indices, g_values, g_shape) 286 # check coo attr 287 compare_coo(coo1, coo2) 288 # check astype 289 compare_coo(py_coo, g_coo) 290 # check to_tuple 291 assert len(py_tuple) == len(g_tuple) 292 for i, _ in enumerate(py_tuple): 293 if isinstance(py_tuple[i], Tensor): 294 assert (py_tuple[i].asnumpy() == g_tuple[i].asnumpy()).all() 295 else: 296 assert py_tuple[i] == g_tuple[i] 297 298 299@pytest.mark.level2 300@pytest.mark.platform_x86_gpu_training 301@pytest.mark.platform_x86_cpu 302@pytest.mark.env_onecard 303def test_coo_bprop(): 304 """ 305 Feature: Test back-propagation with COO-related Ops. 306 Description: Test back-propagation of make_coo, coo.attributes, coo.methods(). 307 Expectation: Success. 308 """ 309 if get_platform() != "linux": 310 return 311 grad_op = ops.GradOperation(get_all=True) 312 indices = Tensor([[0, 1], [1, 2]], dtype=mstype.int32) 313 values = Tensor([-1, 2], dtype=mstype.float32) 314 dense_shape = (3, 4) 315 316 @grad_op 317 @jit 318 def test_coo_tensor(indices, values, dense_shape): 319 coo_tensor = COOTensor(indices, values, dense_shape) 320 return coo_tensor 321 322 @grad_op 323 @jit 324 def test_coo_indices(indices, values, dense_shape): 325 coo_tensor = COOTensor(indices, values, dense_shape) 326 return coo_tensor.indices 327 328 @grad_op 329 @jit 330 def test_coo_values(indices, values, dense_shape): 331 coo_tensor = COOTensor(indices, values, dense_shape) 332 return coo_tensor.values 333 334 @grad_op 335 @jit 336 def test_coo_shape(indices, values, dense_shape): 337 coo_tensor = COOTensor(indices, values, dense_shape) 338 return coo_tensor.shape 339 340 @grad_op 341 @jit 342 def test_coo_cast(indices, values, dense_shape): 343 coo_tensor = COOTensor(indices, values, dense_shape) 344 return coo_tensor.astype(mstype.int32) 345 346 @grad_op 347 @jit 348 def test_coo_dtype(indices, values, dense_shape): 349 coo_tensor = COOTensor(indices, values, dense_shape) 350 return coo_tensor.dtype 351 352 @grad_op 353 @jit 354 def test_coo_to_tuple(indices, values, dense_shape): 355 coo_tensor = COOTensor(indices, values, dense_shape) 356 return coo_tensor.to_tuple() 357 358 @grad_op 359 @jit 360 def test_coo_to_abs(indices, values, dense_shape): 361 coo_tensor = COOTensor(indices, values, dense_shape) 362 return coo_tensor.abs() 363 364 @grad_op 365 @jit 366 def test_coo_to_csr(indices, values, dense_shape): 367 coo_tensor = COOTensor(indices, values, dense_shape) 368 return coo_tensor.to_csr() 369 370 @grad_op 371 @jit 372 def test_coo_to_dense(indices, values, dense_shape): 373 coo_tensor = COOTensor(indices, values, dense_shape) 374 return coo_tensor.to_dense() 375 376 all_zero = (np.zeros(indices.shape, np.int32), np.zeros(values.shape, np.float32)) 377 values_on = (np.zeros(indices.shape, np.int32), np.ones(values.shape, np.float32)) 378 values_absgrad = (np.zeros(indices.shape, np.int32), np.sign(values.asnumpy())) 379 380 compare_res(test_coo_tensor(indices, values, dense_shape), values_on) 381 compare_res(test_coo_indices(indices, values, dense_shape), all_zero) 382 compare_res(test_coo_values(indices, values, dense_shape), values_on) 383 compare_res(test_coo_shape(indices, values, dense_shape), all_zero) 384 compare_res(test_coo_cast(indices, values, dense_shape), values_on) 385 compare_res(test_coo_dtype(indices, values, dense_shape), all_zero) 386 compare_res(test_coo_to_tuple(indices, values, dense_shape), values_on) 387 compare_res(test_coo_to_abs(indices, values, dense_shape), values_absgrad) 388 compare_res(test_coo_to_csr(indices, values, dense_shape), values_on) 389 compare_res(test_coo_to_dense(indices, values, dense_shape), values_on) 390 391 392@pytest.mark.level1 393@pytest.mark.platform_x86_gpu_training 394@pytest.mark.env_onecard 395def test_dense_to_coo(): 396 """ 397 Feature: Test dense tensor to coo methods. 398 Description: Test tensor.to_coo(). 399 Expectation: Success. 400 """ 401 dense_tensor = Tensor([[0, 1, 2, 0], [0, 0, 0, 0], [1, 0, 0, 0]], dtype=mstype.float32) 402 403 def test_to_coo(dense_tensor): 404 return dense_tensor.to_coo() 405 406 coo_tensor = test_to_coo(dense_tensor) 407 coo_tensor_graph = jit(test_to_coo)(dense_tensor) 408 expect = COOTensor(Tensor([[0, 1], [0, 2], [2, 0]], dtype=mstype.int32), 409 Tensor([1, 2, 1], dtype=mstype.float32), 410 (3, 4)) 411 assert isinstance(coo_tensor, COOTensor) 412 assert isinstance(coo_tensor, COOTensor) 413 compare_coo(coo_tensor, expect) 414 compare_coo(coo_tensor_graph, expect) 415 416 417@pytest.mark.level1 418@pytest.mark.platform_x86_gpu_training 419@pytest.mark.platform_x86_cpu 420@pytest.mark.env_onecard 421def test_coo_magic_methods(): 422 """ 423 Feature: Test coo magic methods. 424 Description: Test COOTensor.__neg__, COOTensor.__add__, COOTensor.__sub__, COOTensor.__mul__, COOTensor.__div__. 425 Expectation: Success. 426 """ 427 if get_platform() != "linux": 428 return 429 indices = Tensor([[0, 1], [1, 2]], dtype=mstype.int64) 430 values = Tensor([-1, 2], dtype=mstype.float32) 431 indices = Tensor([[0, 1], [1, 2]], dtype=mstype.int64) 432 values = Tensor([-1, 2], dtype=mstype.float32) 433 shape = (3, 4) 434 dense = Tensor([[0, 1, 2, 0], [0, 0, 2, 0], [1, 0, 0, 0]], dtype=mstype.float32) 435 436 indices_2 = Tensor([[0, 2], [1, 2], [2, 3]], dtype=mstype.int64) 437 values_2 = Tensor([3, -2, 1], dtype=mstype.float32) 438 439 def test_coo_neg(indices, values, shape): 440 coo_tensor = COOTensor(indices, values, shape) 441 return -coo_tensor 442 443 def test_coo_add_coo(indices, indices_2, values, values_2, shape): 444 coo_tensor_1 = COOTensor(indices, values, shape) 445 coo_tensor_2 = COOTensor(indices_2, values_2, shape) 446 return coo_tensor_1 + coo_tensor_2 447 448 def test_coo_add_dense(indices, values, shape, dense): 449 coo_tensor = COOTensor(indices, values, shape) 450 return coo_tensor + dense 451 452 def test_dense_add_coo(indices, values, shape, dense): 453 coo_tensor = COOTensor(indices, values, shape) 454 return dense + coo_tensor 455 456 def test_coo_sub_coo(indices, indices_2, values, values_2, shape): 457 coo_tensor_1 = COOTensor(indices, values, shape) 458 coo_tensor_2 = COOTensor(indices_2, values_2, shape) 459 return coo_tensor_1 - coo_tensor_2 460 461 def test_coo_sub_dense(indices, values, shape, dense): 462 coo_tensor = COOTensor(indices, values, shape) 463 return coo_tensor - dense 464 465 def test_dense_sub_coo(indices, values, shape, dense): 466 coo_tensor = COOTensor(indices, values, shape) 467 return dense - coo_tensor 468 469 def test_coo_mul_dense(indices, values, shape, dense): 470 coo_tensor = COOTensor(indices, values, shape) 471 return coo_tensor * dense 472 473 def test_dense_mul_coo(indices, values, shape, dense): 474 coo_tensor = COOTensor(indices, values, shape) 475 return dense * coo_tensor 476 477 def test_coo_div_dense(indices, values, shape, dense): 478 coo_tensor = COOTensor(indices, values, shape) 479 return coo_tensor / dense 480 481 neg_expect = COOTensor(indices, Tensor([1, -2], mstype.float32), shape) 482 neg_output = test_coo_neg(indices, values, shape) 483 compare_coo(neg_output, neg_expect) 484 neg_output = jit(test_coo_neg)(indices, values, shape) 485 compare_coo(neg_output, neg_expect) 486 487 coo_add_coo_expect = COOTensor(Tensor([[0, 1], [0, 2], [1, 2], [2, 3]], mstype.int64), 488 Tensor([-1, 3, 0, 1], mstype.float32), shape) 489 coo_add_coo_output = test_coo_add_coo(indices, indices_2, values, values_2, shape) 490 compare_coo(coo_add_coo_output, coo_add_coo_expect) 491 coo_add_coo_output = jit(test_coo_add_coo)(indices, indices_2, values, values_2, shape) 492 compare_coo(coo_add_coo_output, coo_add_coo_expect) 493 494 coo_add_dense_expect = np.array([[0, 0, 2, 0], [0, 0, 4, 0,], [1, 0, 0, 0,]], np.int64) 495 coo_add_dense_output = test_coo_add_dense(indices, values, shape, dense) 496 assert np.allclose(coo_add_dense_expect, coo_add_dense_output.asnumpy()) 497 coo_add_dense_output = jit(test_coo_add_dense)(indices, values, shape, dense) 498 assert np.allclose(coo_add_dense_expect, coo_add_dense_output.asnumpy()) 499 500 dense_add_coo_output = test_dense_add_coo(indices, values, shape, dense) 501 assert np.allclose(coo_add_dense_expect, dense_add_coo_output.asnumpy()) 502 dense_add_coo_output = jit(test_dense_add_coo)(indices, values, shape, dense) 503 assert np.allclose(coo_add_dense_expect, dense_add_coo_output.asnumpy()) 504 505 coo_sub_coo_expect = COOTensor(Tensor([[0, 1], [0, 2], [1, 2], [2, 3]], mstype.int64), 506 Tensor([-1, -3, 4, -1], mstype.float32), shape) 507 coo_sub_coo_output = test_coo_sub_coo(indices, indices_2, values, values_2, shape) 508 compare_coo(coo_sub_coo_output, coo_sub_coo_expect) 509 coo_sub_coo_output = jit(test_coo_sub_coo)(indices, indices_2, values, values_2, shape) 510 compare_coo(coo_sub_coo_output, coo_sub_coo_expect) 511 512 coo_sub_dense_expect = np.array([[0, -2, -2, 0], [0, 0, 0, 0], [-1, 0, 0, 0]], np.int32) 513 coo_sub_dense_output = test_coo_sub_dense(indices, values, shape, dense) 514 assert np.allclose(coo_sub_dense_expect, coo_sub_dense_output.asnumpy()) 515 coo_sub_dense_output = jit(test_coo_sub_dense)(indices, values, shape, dense) 516 assert np.allclose(coo_sub_dense_expect, coo_sub_dense_output.asnumpy()) 517 518 dense_sub_coo_expect = np.array([[0, 2, 2, 0], [0, 0, 0, 0], [1, 0, 0, 0]], np.int64) 519 dense_sub_coo_output = test_dense_sub_coo(indices, values, shape, dense) 520 assert np.allclose(dense_sub_coo_expect, dense_sub_coo_output.asnumpy()) 521 dense_sub_coo_output = jit(test_dense_sub_coo)(indices, values, shape, dense) 522 assert np.allclose(dense_sub_coo_expect, dense_sub_coo_output.asnumpy()) 523 524 coo_mul_dense_expect = COOTensor(indices, Tensor([-1, 4], mstype.float32), shape) 525 coo_mul_dense_output = test_coo_mul_dense(indices, values, shape, dense) 526 compare_coo(coo_mul_dense_output, coo_mul_dense_expect) 527 coo_mul_dense_output = jit(test_coo_mul_dense)(indices, values, shape, dense) 528 compare_coo(coo_mul_dense_output, coo_mul_dense_expect) 529 530 dense_mul_coo_output = test_dense_mul_coo(indices, values, shape, dense) 531 compare_coo(dense_mul_coo_output, coo_mul_dense_expect) 532 dense_mul_coo_output = jit(test_dense_mul_coo)(indices, values, shape, dense) 533 compare_coo(dense_mul_coo_output, coo_mul_dense_expect) 534 535 coo_div_dense_expect = COOTensor(indices, Tensor([-1, 1], mstype.float32), shape) 536 coo_div_dense_output = test_coo_div_dense(indices, values, shape, dense) 537 compare_coo(coo_div_dense_output, coo_div_dense_expect) 538 coo_div_dense_output = jit(test_coo_div_dense)(indices, values, shape, dense) 539 compare_coo(coo_div_dense_output, coo_div_dense_expect) 540