| 308 | |
| 309 | |
| 310 | def drop_node(edge_index: torch.Tensor, edge_weight: Optional[torch.Tensor] = None, keep_prob: float = 0.5) -> (torch.Tensor, Optional[torch.Tensor]): |
| 311 | num_nodes = edge_index.max().item() + 1 |
| 312 | probs = torch.tensor([keep_prob for _ in range(num_nodes)]) |
| 313 | dist = Bernoulli(probs) |
| 314 | |
| 315 | subset = dist.sample().to(torch.bool).to(edge_index.device) |
| 316 | edge_index, edge_weight = subgraph(subset, edge_index, edge_weight) |
| 317 | |
| 318 | return edge_index, edge_weight |
| 319 | |
| 320 | |
| 321 | def random_walk_subgraph(edge_index: torch.LongTensor, edge_weight: Optional[torch.FloatTensor] = None, batch_size: int = 1000, length: int = 10): |