def edge_sampler(self, g):
        """Node ID sampler for random edge sampler"""
        src, dst = g.edges()
        if self.cache and self.prob is not None:
            prob = self.prob
        else:
            in_deg = g.in_degrees().float().clamp(min=1)
            out_deg = g.out_degrees().float().clamp(min=1)
            # We can reduce the sample space by half if graphs are always symmetric.
            prob = 1.0 / in_deg[dst.long()] + 1.0 / out_deg[src.long()]
            prob /= prob.sum()
            if self.cache:
                self.prob = prob
        sampled_edges = torch.unique(
            choice(len(prob), size=self.budget, prob=prob)
        )
        sampled_nodes = torch.cat([src[sampled_edges], dst[sampled_edges]])
        return sampled_nodes.unique().type(g.idtype)
Node ID sampler for random edge sampler
python/dgl/dataloading/graphsaint.py