Samples a list of blocks, as well as a subgraph containing the sampled edges from the original graph. If :attr:`negative_sampler` is given, also returns another graph containing the negative pairs as edges.
(self, g, seed_edges)
| 451 | return result |
| 452 | |
| 453 | def sample(self, g, seed_edges): # pylint: disable=arguments-differ |
| 454 | """Samples a list of blocks, as well as a subgraph containing the sampled |
| 455 | edges from the original graph. |
| 456 | |
| 457 | If :attr:`negative_sampler` is given, also returns another graph containing the |
| 458 | negative pairs as edges. |
| 459 | """ |
| 460 | if isinstance(seed_edges, Mapping): |
| 461 | seed_edges = { |
| 462 | g.to_canonical_etype(k): v for k, v in seed_edges.items() |
| 463 | } |
| 464 | exclude = self.exclude |
| 465 | pair_graph = g.edge_subgraph( |
| 466 | seed_edges, relabel_nodes=False, output_device=self.output_device |
| 467 | ) |
| 468 | eids = pair_graph.edata[EID] |
| 469 | |
| 470 | if self.negative_sampler is not None: |
| 471 | neg_graph = self._build_neg_graph(g, seed_edges) |
| 472 | pair_graph, neg_graph = compact_graphs([pair_graph, neg_graph]) |
| 473 | else: |
| 474 | pair_graph = compact_graphs(pair_graph) |
| 475 | |
| 476 | pair_graph.edata[EID] = eids |
| 477 | seed_nodes = pair_graph.ndata[NID] |
| 478 | |
| 479 | exclude_eids = find_exclude_eids( |
| 480 | g, |
| 481 | seed_edges, |
| 482 | exclude, |
| 483 | self.reverse_eids, |
| 484 | self.reverse_etypes, |
| 485 | self.output_device, |
| 486 | ) |
| 487 | |
| 488 | input_nodes, _, blocks = self.sampler.sample( |
| 489 | g, seed_nodes, exclude_eids |
| 490 | ) |
| 491 | |
| 492 | if self.negative_sampler is None: |
| 493 | return self.assign_lazy_features((input_nodes, pair_graph, blocks)) |
| 494 | else: |
| 495 | return self.assign_lazy_features( |
| 496 | (input_nodes, pair_graph, neg_graph, blocks) |
| 497 | ) |
| 498 | |
| 499 | |
| 500 | def as_edge_prediction_sampler( |
nothing calls this directly
no test coverage detected