Returns a sub-matrix selected according to the given range index. Parameters ---------- dim : int The dim to select from matrix, should be 0 or 1. `dim = 0` for rowwise selection and `dim = 1` for columnwise selection. index : slice
(self, dim: int, index: slice)
| 531 | raise TypeError(f"{type(index).__name__} is unsupported input type.") |
| 532 | |
| 533 | def range_select(self, dim: int, index: slice): |
| 534 | """Returns a sub-matrix selected according to the given range index. |
| 535 | |
| 536 | Parameters |
| 537 | ---------- |
| 538 | dim : int |
| 539 | The dim to select from matrix, should be 0 or 1. `dim = 0` for |
| 540 | rowwise selection and `dim = 1` for columnwise selection. |
| 541 | index : slice |
| 542 | The selection slice indicates ID index from the `dim` should |
| 543 | be chosen from the matrix. |
| 544 | |
| 545 | The function does not support autograd. |
| 546 | |
| 547 | Returns |
| 548 | ------- |
| 549 | SparseMatrix |
| 550 | The sub-matrix which contains selected rows or columns. |
| 551 | |
| 552 | Examples |
| 553 | -------- |
| 554 | |
| 555 | >>> indices = torch.tensor([0, 1, 1, 2, 3, 4], [0, 2, 4, 3, 5, 0]]) |
| 556 | >>> val = torch.tensor([0, 1, 2, 3, 4, 5]) |
| 557 | >>> A = dglsp.spmatrix(indices, val) |
| 558 | |
| 559 | Case 1: Select rows with given slice object. |
| 560 | |
| 561 | >>> A.range_select(0, slice(1, 3)) |
| 562 | SparseMatrix(indices=tensor([[0, 0, 1], |
| 563 | [2, 4, 3]]), |
| 564 | values=tensor([1, 2, 3]), |
| 565 | shape=(2, 6), nnz=3) |
| 566 | |
| 567 | Case 2: Select columns with given slice object. |
| 568 | |
| 569 | >>> A.range_select(1, slice(3, 6)) |
| 570 | SparseMatrix(indices=tensor([[2, 1, 3], |
| 571 | [0, 1, 2]]), |
| 572 | values=tensor([3, 2, 4]), |
| 573 | shape=(5, 3), nnz=3) |
| 574 | """ |
| 575 | if dim not in (0, 1): |
| 576 | raise ValueError("The selection dimension should be 0 or 1.") |
| 577 | if isinstance(index, slice): |
| 578 | if index.step not in (None, 1): |
| 579 | raise NotImplementedError( |
| 580 | "Slice with step other than 1 are not supported yet." |
| 581 | ) |
| 582 | start = 0 if index.start is None else index.start |
| 583 | end = index.stop |
| 584 | return SparseMatrix( |
| 585 | self.c_sparse_matrix.range_select(dim, start, end) |
| 586 | ) |
| 587 | raise TypeError(f"{type(index).__name__} is unsupported input type.") |
| 588 | |
| 589 | def sample( |
| 590 | self, |