Relabel the input ids to continuous ids that starts from zero. Ids are assigned new ids according to their ascending order. Examples -------- >>> x = [1, 5, 3, 6] >>> n2o, o2n = build_relabel_map(x) >>> n2o [1, 3, 5, 6] >>> o2n [n/a, 0, n/a, 1, n/a, 2, 3] "
(x)
| 799 | |
| 800 | |
| 801 | def relabel(x): |
| 802 | """Relabel the input ids to continuous ids that starts from zero. |
| 803 | |
| 804 | Ids are assigned new ids according to their ascending order. |
| 805 | |
| 806 | Examples |
| 807 | -------- |
| 808 | >>> x = [1, 5, 3, 6] |
| 809 | >>> n2o, o2n = build_relabel_map(x) |
| 810 | >>> n2o |
| 811 | [1, 3, 5, 6] |
| 812 | >>> o2n |
| 813 | [n/a, 0, n/a, 1, n/a, 2, 3] |
| 814 | |
| 815 | "n/a" will be filled with 0 |
| 816 | |
| 817 | Parameters |
| 818 | ---------- |
| 819 | x : Tensor |
| 820 | ID tensor. |
| 821 | |
| 822 | Returns |
| 823 | ------- |
| 824 | new_to_old : Tensor |
| 825 | The mapping from new id to old id. |
| 826 | old_to_new : Tensor |
| 827 | The mapping from old id to new id. It is a vector of length MAX(x). |
| 828 | One can use advanced indexing to convert an old id tensor to a |
| 829 | new id tensor: new_id = old_to_new[old_id] |
| 830 | """ |
| 831 | unique_x = F.unique(x) |
| 832 | map_len = F.as_scalar(F.max(unique_x, dim=0)) + 1 |
| 833 | ctx = F.context(x) |
| 834 | dtype = F.dtype(x) |
| 835 | old_to_new = F.zeros((map_len,), dtype=dtype, ctx=ctx) |
| 836 | old_to_new = F.scatter_row( |
| 837 | old_to_new, unique_x, F.copy_to(F.arange(0, len(unique_x), dtype), ctx) |
| 838 | ) |
| 839 | return unique_x, old_to_new |
| 840 | |
| 841 | |
| 842 | def extract_node_subframes(graph, nodes_or_device, store_ids=True): |