def _sample_neighbors_graphbolt(
    g,
    gpb,
    nodes,
    fanout,
    edge_dir="in",
    prob=None,
    exclude_edges=None,
    replace=False,
):
    """Sample from local partition via graphbolt.

    The input nodes use global IDs. We need to map the global node IDs to local
    node IDs, perform sampling and map the sampled results to the global IDs
    space again. The sampled results are stored in three vectors that store
    source nodes, destination nodes, etype IDs and edge IDs.

    Parameters
    ----------
    g : FusedCSCSamplingGraph
        The local partition.
    gpb : GraphPartitionBook
        The graph partition book.
    nodes : tensor
        The nodes to sample neighbors from.
    fanout : tensor or int
        The number of edges to be sampled for each node.
    edge_dir : str, optional
        Determines whether to sample inbound or outbound edges.
    prob : tensor, optional
        The probability associated with each neighboring edge of a node.
    exclude_edges : tensor, optional
        The edges to exclude when sampling.
    replace : bool, optional
        If True, sample with replacement.

    Returns
    -------
    tensor
        The source node ID array.
    tensor
        The destination node ID array.
    tensor
        The edge ID array.
    tensor
        The edge type ID array.
    """
    assert (
        edge_dir == "in"
    ), f"GraphBolt only supports inbound edge sampling but got {edge_dir}."
    assert exclude_edges is None, "GraphBolt does not support excluding edges."

    # 1. Map global node IDs to local node IDs.
    nodes = gpb.nid2localnid(nodes, gpb.partid)
    # Local partition may be saved in torch.int32 even though the global graph
    # is in torch.int64.
    nodes = nodes.to(dtype=g.indices.dtype)

    # 2. Perform sampling.
    probs_or_mask = None
    if prob is not None:
        probs_or_mask = g.edge_attributes[prob]
    # Sanity checks.
    assert isinstance(
        g, gb.FusedCSCSamplingGraph
    ), "Expect a FusedCSCSamplingGraph."
    assert isinstance(nodes, torch.Tensor), "Expect a tensor of nodes."
    if isinstance(fanout, int):
        fanout = torch.LongTensor([fanout])
    assert isinstance(fanout, torch.Tensor), "Expect a tensor of fanout."

    subgraph = g._sample_neighbors(
        nodes,
        None,
        fanout,
        replace=replace,
        probs_or_mask=probs_or_mask,
    )

    # 3. Map local node IDs to global node IDs.
    local_src = subgraph.indices
    local_dst = gb.expand_indptr(
        subgraph.indptr,
        dtype=local_src.dtype,
        node_ids=subgraph.original_column_node_ids,
        output_size=local_src.shape[0],
    )
    global_nid_mapping = g.node_attributes[NID]
    global_src = global_nid_mapping[local_src]
    global_dst = global_nid_mapping[local_dst]

    global_eids = None
    if g.edge_attributes is not None and EID in g.edge_attributes:
        global_eids = g.edge_attributes[EID][subgraph.original_edge_ids]
    return LocalSampledGraph(
        global_src, global_dst, global_eids, subgraph.type_per_edge
    )
Sample from local partition via graphbolt.

    The input nodes use global IDs. We need to map the global node IDs to local
    node IDs, perform sampling and map the sampled results to the global IDs
    space again. The sampled results are stored in three vectors that store
    source nodes, destination nodes, etype IDs and edge IDs.

    Parameters
    ----------
    g : FusedCSCSamplingGraph
        The local partition.
    gpb : GraphPartitionBook
        The graph partition book.
    nodes : tensor
        The nodes to sample neighbors from.
    fanout : tensor or int
        The number of edges to be sampled for each node.
    edge_dir : str, optional
        Determines whether to sample inbound or outbound edges.
    prob : tensor, optional
        The probability associated with each neighboring edge of a node.
    exclude_edges : tensor, optional
        The edges to exclude when sampling.
    replace : bool, optional
        If True, sample with replacement.

    Returns
    -------
    tensor
        The source node ID array.
    tensor
        The destination node ID array.
    tensor
        The edge ID array.
    tensor
        The edge type ID array.
    
python/dgl/distributed/graph_services.py