def sample_etype_neighbors(
    g,
    nodes,
    fanout,
    edge_dir="in",
    prob=None,
    exclude_edges=None,
    replace=False,
    etype_sorted=True,
    use_graphbolt=False,
):
    """Sample from the neighbors of the given nodes from a distributed graph.

    For each node, a number of inbound (or outbound when ``edge_dir == 'out'``) edges
    will be randomly chosen.  The returned graph will contain all the nodes in the
    original graph, but only the sampled edges.

    Node/edge features are not preserved. The original IDs of
    the sampled edges are stored as the `dgl.EID` feature in the returned graph.

    This function assumes the input is a homogeneous ``DGLGraph`` with the edges
    ordered by their edge types. The sampled subgraph is also
    stored in the homogeneous graph format. That is, all nodes and edges are assigned
    with unique IDs (in contrast, we typically use a type name and a node/edge ID to
    identify a node or an edge in ``DGLGraph``). We refer to this type of IDs
    as *homogeneous ID*.
    Users can use :func:`dgl.distributed.GraphPartitionBook.map_to_per_ntype`
    and :func:`dgl.distributed.GraphPartitionBook.map_to_per_etype`
    to identify their node/edge types and node/edge IDs of that type.

    Parameters
    ----------
    g : DistGraph
        The distributed graph..
    nodes : tensor or dict
        Node IDs to sample neighbors from. If it's a dict, it should contain only
        one key-value pair to make this API consistent with dgl.sampling.sample_neighbors.
    fanout : int or dict[etype, int]
        The number of edges to be sampled for each node per edge type.  If an integer
        is given, DGL assumes that the same fanout is applied to every edge type.

        If -1 is given, all of the neighbors will be selected.
    edge_dir : str, optional
        Determines whether to sample inbound or outbound edges.

        Can take either ``in`` for inbound edges or ``out`` for outbound edges.
    prob : str, optional
        Feature name used as the (unnormalized) probabilities associated with each
        neighboring edge of a node.  The feature must have only one element for each
        edge.

        The features must be non-negative floats, and the sum of the features of
        inbound/outbound edges for every node must be positive (though they don't have
        to sum up to one).  Otherwise, the result will be undefined.
    exclude_edges : tensor, optional
        The edges to exclude when sampling. Homogeneous edge IDs are used.
    replace : bool, optional
        If True, sample with replacement.

        When sampling with replacement, the sampled subgraph could have parallel edges.

        For sampling without replacement, if fanout > the number of neighbors, all the
        neighbors are sampled. If fanout == -1, all neighbors are collected.
    etype_sorted : bool, optional
        Indicates whether etypes are sorted.
    use_graphbolt : bool, optional
        Whether to use GraphBolt for sampling.

    Returns
    -------
    DGLGraph
        A sampled subgraph containing only the sampled neighboring edges.  It is on CPU.
    """
    if isinstance(fanout, int):
        fanout = F.full_1d(len(g.canonical_etypes), fanout, F.int64, F.cpu())
    else:
        etype_ids = {etype: i for i, etype in enumerate(g.canonical_etypes)}
        fanout_array = [None] * len(g.canonical_etypes)
        for etype, v in fanout.items():
            c_etype = g.to_canonical_etype(etype)
            fanout_array[etype_ids[c_etype]] = v
        assert all(v is not None for v in fanout_array), (
            "Not all etypes have valid fanout. Please make sure passed-in "
            "fanout in dict includes all the etypes in graph. Passed-in "
            f"fanout: {fanout}, graph etypes: {g.canonical_etypes}."
        )
        fanout = F.tensor(fanout_array, dtype=F.int64)

    gpb = g.get_partition_book()
    if isinstance(nodes, dict):
        homo_nids = []
        for ntype in nodes.keys():
            assert (
                ntype in g.ntypes
            ), "The sampled node type {} does not exist in the input graph".format(
                ntype
            )
            if F.is_tensor(nodes[ntype]):
                typed_nodes = nodes[ntype]
            else:
                typed_nodes = toindex(nodes[ntype]).tousertensor()
            homo_nids.append(gpb.map_to_homo_nid(typed_nodes, ntype))
        nodes = F.cat(homo_nids, 0)

    def issue_remote_req(node_ids):
        if prob is not None and (not use_graphbolt):
            # See NOTE 1
            _prob = [
                (
                    # NOTE (BarclayII)
                    # Currently DistGraph.edges[] does not accept canonical etype.
                    g.edges[etype].data[prob].kvstore_key
                    if prob in g.edges[etype].data
                    else ""
                )
                for etype in g.canonical_etypes
            ]
        else:
            _prob = prob
        return SamplingRequestEtype(
            node_ids,
            fanout,
            edge_dir=edge_dir,
            prob=_prob,
            exclude_edges=None,
            replace=replace,
            etype_sorted=etype_sorted,
            use_graphbolt=use_graphbolt,
        )

    def local_access(local_g, partition_book, local_nids):
        etype_offset = gpb.local_etype_offset
        # See NOTE 1
        if prob is not None and (not use_graphbolt):
            _prob = [
                (
                    g.edges[etype].data[prob].local_partition
                    if prob in g.edges[etype].data
                    else None
                )
                for etype in g.canonical_etypes
            ]
        else:
            _prob = prob
        return _sample_etype_neighbors(
            use_graphbolt,
            local_g,
            partition_book,
            local_nids,
            fanout,
            edge_dir=edge_dir,
            prob=_prob,
            exclude_edges=None,
            replace=replace,
            etype_offset=etype_offset,
            etype_sorted=etype_sorted,
        )

    frontier = _distributed_access(
        g, nodes, issue_remote_req, local_access, exclude_edges=exclude_edges
    )
    if not gpb.is_homogeneous:
        return _frontier_to_heterogeneous_graph(g, frontier, gpb)
    else:
        return frontier
Sample from the neighbors of the given nodes from a distributed graph.

    For each node, a number of inbound (or outbound when ``edge_dir == 'out'``) edges
    will be randomly chosen.  The returned graph will contain all the nodes in the
    original graph, but only the sampled edges.

    Node/edge features are not preserved. The original IDs of
    the sampled edges are stored as the `dgl.EID` feature in the returned graph.

    This function assumes the input is a homogeneous ``DGLGraph`` with the edges
    ordered by their edge types. The sampled subgraph is also
    stored in the homogeneous graph format. That is, all nodes and edges are assigned
    with unique IDs (in contrast, we typically use a type name and a node/edge ID to
    identify a node or an edge in ``DGLGraph``). We refer to this type of IDs
    as *homogeneous ID*.
    Users can use :func:`dgl.distributed.GraphPartitionBook.map_to_per_ntype`
    and :func:`dgl.distributed.GraphPartitionBook.map_to_per_etype`
    to identify their node/edge types and node/edge IDs of that type.

    Parameters
    ----------
    g : DistGraph
        The distributed graph..
    nodes : tensor or dict
        Node IDs to sample neighbors from. If it's a dict, it should contain only
        one key-value pair to make this API consistent with dgl.sampling.sample_neighbors.
    fanout : int or dict[etype, int]
        The number of edges to be sampled for each node per edge type.  If an integer
        is given, DGL assumes that the same fanout is applied to every edge type.

        If -1 is given, all of the neighbors will be selected.
    edge_dir : str, optional
        Determines whether to sample inbound or outbound edges.

        Can take either ``in`` for inbound edges or ``out`` for outbound edges.
    prob : str, optional
        Feature name used as the (unnormalized) probabilities associated with each
        neighboring edge of a node.  The feature must have only one element for each
        edge.

        The features must be non-negative floats, and the sum of the features of
        inbound/outbound edges for every node must be positive (though they don't have
        to sum up to one).  Otherwise, the result will be undefined.
    exclude_edges : tensor, optional
        The edges to exclude when sampling. Homogeneous edge IDs are used.
    replace : bool, optional
        If True, sample with replacement.

        When sampling with replacement, the sampled subgraph could have parallel edges.

        For sampling without replacement, if fanout > the number of neighbors, all the
        neighbors are sampled. If fanout == -1, all neighbors are collected.
    etype_sorted : bool, optional
        Indicates whether etypes are sorted.
    use_graphbolt : bool, optional
        Whether to use GraphBolt for sampling.

    Returns
    -------
    DGLGraph
        A sampled subgraph containing only the sampled neighboring edges.  It is on CPU.
    
python/dgl/distributed/graph_services.py