Produce a scheme that maps the given indices to separate partitions and the counts of how many indices are in each partition. Parameters ---------- idxs: torch.Tensor. A tensor with shape (`num_indices`,), representing global indices. Return
(self, idxs)
| 593 | ) |
| 594 | |
| 595 | def generate_permutation(self, idxs): |
| 596 | """Produce a scheme that maps the given indices to separate partitions |
| 597 | and the counts of how many indices are in each partition. |
| 598 | |
| 599 | |
| 600 | Parameters |
| 601 | ---------- |
| 602 | idxs: torch.Tensor. |
| 603 | A tensor with shape (`num_indices`,), representing global indices. |
| 604 | |
| 605 | Return |
| 606 | ------ |
| 607 | torch.Tensor. |
| 608 | A tensor with shape (`num_indices`,), representing the permutation |
| 609 | to re-order the indices by partition. |
| 610 | torch.Tensor. |
| 611 | A tensor with shape (`num_partition`,), representing the number of |
| 612 | indices per partition. |
| 613 | |
| 614 | Examples |
| 615 | -------- |
| 616 | |
| 617 | >>> import torch |
| 618 | >>> from dgl.partition import NDArrayPartition |
| 619 | >>> part = NDArrayPartition(10, 2, mode="remainder") |
| 620 | >>> idx = torch.tensor([0, 2, 4, 5, 8, 8, 9], device="cuda:0") |
| 621 | >>> perm, splits_sum = part.generate_permutation(idx) |
| 622 | >>> perm |
| 623 | tensor([0, 1, 2, 4, 5, 3, 6], device='cuda:0') |
| 624 | >>> splits_sum |
| 625 | tensor([5, 2], device='cuda:0') |
| 626 | """ |
| 627 | ret = _CAPI_DGLNDArrayPartitionGeneratePermutation( |
| 628 | self._partition, F.zerocopy_to_dgl_ndarray(idxs) |
| 629 | ) |
| 630 | return F.zerocopy_from_dgl_ndarray(ret(0)), F.zerocopy_from_dgl_ndarray( |
| 631 | ret(1) |
| 632 | ) |
| 633 | |
| 634 | |
| 635 | _init_api("dgl.partition") |
no outgoing calls