def sample(
        self, g, indices, exclude_eids=None
    ):  # pylint: disable=arguments-differ
        """Sampling function.

        Parameters
        ----------
        g : DGLGraph
            The graph to sample from.
        indices : Tensor or dict[str, Tensor]
            Nodes which induce the subgraph.
        exclude_eids : Tensor or dict[etype, Tensor], optional
            The edges to exclude from the sampled subgraph.

        Returns
        -------
        input_nodes : Tensor or dict[str, Tensor]
            The node IDs inducing the subgraph.
        output_nodes : Tensor or dict[str, Tensor]
            The node IDs that are sampled in this minibatch.
        subg : DGLGraph
            The subgraph itself.
        """

        # Define empty dictionary to store reached nodes.
        output_nodes = indices
        all_reached_nodes = [indices]

        # Iterate over fanout.
        for fanout in reversed(self.fanouts):

            # Sample frontier.
            frontier = g.sample_neighbors(
                indices,
                fanout,
                output_device=self.output_device,
                replace=self.replace,
                prob=self.prob,
                exclude_edges=exclude_eids,
            )

            # Get reached nodes.
            curr_reached = defaultdict(list)
            for c_etype in frontier.canonical_etypes:
                (src_type, _, _) = c_etype
                src, _ = frontier.edges(etype=c_etype)
                curr_reached[src_type].append(src)

            # De-duplication.
            curr_reached = {
                ntype: torch.unique(torch.cat(srcs))
                for ntype, srcs in curr_reached.items()
            }

            # Generate type sampling probabilties.
            type_count = {
                node_type: indices.shape[0]
                for node_type, indices in curr_reached.items()
            }
            total_count = sum(type_count.values())
            probs = {
                node_type: count / total_count
                for node_type, count in type_count.items()
            }

            # Upsample rare node types.
            if self.upsample_rare_types:

                # Take scaled square root of probabilities.
                prob_dist = list(probs.values())
                prob_dist = np.sqrt(prob_dist)
                prob_dist = prob_dist / prob_dist.sum()

                # Update probabilities.
                probs = {
                    node_type: prob_dist[i]
                    for i, node_type in enumerate(probs.keys())
                }

            # Generate node counts per type.
            n_per_type = {
                node_type: int(self.fixed_k * prob)
                for node_type, prob in probs.items()
            }
            remainder = self.fixed_k - sum(n_per_type.values())
            for _ in range(remainder):
                node_type = np.random.choice(
                    list(probs.keys()), p=list(probs.values())
                )
                n_per_type[node_type] += 1

            # Downsample nodes.
            curr_reached_k = {}
            for node_type, node_ids in curr_reached.items():

                # Get number of total nodes and number to sample.
                num_nodes = node_ids.shape[0]
                n_to_sample = min(num_nodes, n_per_type[node_type])

                # Downsample nodes of current type.
                random_indices = torch.randperm(num_nodes)[:n_to_sample]
                curr_reached_k[node_type] = node_ids[random_indices]

            # Update seed nodes.
            indices = curr_reached_k
            all_reached_nodes.append(curr_reached_k)

        # Merge all reached nodes before sending to `DGLGraph.subgraph`.
        merged_nodes = {}
        for ntype in g.ntypes:
            merged_nodes[ntype] = torch.unique(
                torch.cat(
                    [reached.get(ntype, []) for reached in all_reached_nodes]
                )
            )
        subg = g.subgraph(
            merged_nodes, relabel_nodes=True, output_device=self.output_device
        )

        if exclude_eids is not None:
            subg = EidExcluder(exclude_eids)(subg)

        set_node_lazy_features(subg, self.prefetch_node_feats)
        set_edge_lazy_features(subg, self.prefetch_edge_feats)

        return indices, output_nodes, subg
Sampling function.

        Parameters
        ----------
        g : DGLGraph
            The graph to sample from.
        indices : Tensor or dict[str, Tensor]
            Nodes which induce the subgraph.
        exclude_eids : Tensor or dict[etype, Tensor], optional
            The edges to exclude from the sampled subgraph.

        Returns
        -------
        input_nodes : Tensor or dict[str, Tensor]
            The node IDs inducing the subgraph.
        output_nodes : Tensor or dict[str, Tensor]
            The node IDs that are sampled in this minibatch.
        subg : DGLGraph
            The subgraph itself.
        
python/dgl/dataloading/capped_neighbor_sampler.py