def node_sampler(self, g):
        """Node ID sampler for random node sampler"""
        # Alternatively, this can be realized by uniformly sampling an edge subset,
        # and then take the src node of the sampled edges. However, the number of edges
        # is typically much larger than the number of nodes.
        if self.cache and self.prob is not None:
            prob = self.prob
        else:
            prob = g.out_degrees().float().clamp(min=1)
            if self.cache:
                self.prob = prob
        return (
            torch.multinomial(prob, num_samples=self.budget, replacement=True)
            .unique()
            .type(g.idtype)
        )
Node ID sampler for random node sampler
python/dgl/dataloading/graphsaint.py