| 589 | |
| 590 | |
| 591 | def test_print(): |
| 592 | ctx = F.ctx() |
| 593 | |
| 594 | # basic |
| 595 | row = torch.tensor([1, 1, 3]).to(ctx) |
| 596 | col = torch.tensor([2, 1, 3]).to(ctx) |
| 597 | val = torch.tensor([1.0, 1.0, 2.0]).to(ctx) |
| 598 | A = from_coo(row, col, val) |
| 599 | expected = ( |
| 600 | str( |
| 601 | """SparseMatrix(indices=tensor([[1, 1, 3], |
| 602 | [2, 1, 3]]), |
| 603 | values=tensor([1., 1., 2.]), |
| 604 | shape=(4, 4), nnz=3)""" |
| 605 | ) |
| 606 | if str(ctx) == "cpu" |
| 607 | else str( |
| 608 | """SparseMatrix(indices=tensor([[1, 1, 3], |
| 609 | [2, 1, 3]], device='cuda:0'), |
| 610 | values=tensor([1., 1., 2.], device='cuda:0'), |
| 611 | shape=(4, 4), nnz=3)""" |
| 612 | ) |
| 613 | ) |
| 614 | assert str(A) == expected, print(A, expected) |
| 615 | |
| 616 | # vector-shape non zero |
| 617 | row = torch.tensor([1, 1, 3]).to(ctx) |
| 618 | col = torch.tensor([2, 1, 3]).to(ctx) |
| 619 | val = torch.tensor( |
| 620 | [[1.3080, 1.5984], [-0.4126, 0.7250], [-0.5416, -0.7022]] |
| 621 | ).to(ctx) |
| 622 | A = from_coo(row, col, val) |
| 623 | expected = ( |
| 624 | str( |
| 625 | """SparseMatrix(indices=tensor([[1, 1, 3], |
| 626 | [2, 1, 3]]), |
| 627 | values=tensor([[ 1.3080, 1.5984], |
| 628 | [-0.4126, 0.7250], |
| 629 | [-0.5416, -0.7022]]), |
| 630 | shape=(4, 4), nnz=3, val_size=(2,))""" |
| 631 | ) |
| 632 | if str(ctx) == "cpu" |
| 633 | else str( |
| 634 | """SparseMatrix(indices=tensor([[1, 1, 3], |
| 635 | [2, 1, 3]], device='cuda:0'), |
| 636 | values=tensor([[ 1.3080, 1.5984], |
| 637 | [-0.4126, 0.7250], |
| 638 | [-0.5416, -0.7022]], device='cuda:0'), |
| 639 | shape=(4, 4), nnz=3, val_size=(2,))""" |
| 640 | ) |
| 641 | ) |
| 642 | assert str(A) == expected, print(A, expected) |
| 643 | |
| 644 | |
| 645 | @unittest.skipIf( |