Sampling function. Parameters ---------- g : DGLGraph The graph to sample from. partition_ids : Tensor A 1-D integer tensor of partition IDs. Returns ------- DGLGraph The sampled subgraph.
(self, g, partition_ids)
| 124 | self.output_device = output_device |
| 125 | |
| 126 | def sample(self, g, partition_ids): # pylint: disable=arguments-differ |
| 127 | """Sampling function. |
| 128 | |
| 129 | Parameters |
| 130 | ---------- |
| 131 | g : DGLGraph |
| 132 | The graph to sample from. |
| 133 | partition_ids : Tensor |
| 134 | A 1-D integer tensor of partition IDs. |
| 135 | |
| 136 | Returns |
| 137 | ------- |
| 138 | DGLGraph |
| 139 | The sampled subgraph. |
| 140 | """ |
| 141 | node_ids = F.cat( |
| 142 | [ |
| 143 | self.partition_node_ids[ |
| 144 | self.partition_offset[i] : self.partition_offset[i + 1] |
| 145 | ] |
| 146 | for i in F.asnumpy(partition_ids) |
| 147 | ], |
| 148 | 0, |
| 149 | ) |
| 150 | sg = g.subgraph( |
| 151 | node_ids, relabel_nodes=True, output_device=self.output_device |
| 152 | ) |
| 153 | set_node_lazy_features(sg, self.prefetch_ndata) |
| 154 | set_edge_lazy_features(sg, self.prefetch_edata) |
| 155 | return sg |
nothing calls this directly
no test coverage detected