def _exclude_homo_edges(
    edges: Tuple[torch.Tensor, torch.Tensor],
    edges_to_exclude: torch.Tensor,
    assume_num_node_within_int32: bool,
    async_op: bool,
):
    """Return the indices of edges to be included."""
    if assume_num_node_within_int32:
        val = edges[0].long() << 32 | edges[1].long()
        edges_to_exclude_trans = edges_to_exclude.T
        val_to_exclude = (
            edges_to_exclude_trans[0].long() << 32
            | edges_to_exclude_trans[1].long()
        )
    else:
        # TODO: Add support for value > int32.
        raise NotImplementedError(
            "Values out of range int32 are not supported yet"
        )
    if async_op:
        return torch.ops.graphbolt.is_not_in_index_async(val, val_to_exclude)
    else:
        mask = ~isin(val, val_to_exclude)
        return torch.nonzero(mask, as_tuple=True)[0]
Return the indices of edges to be included.
python/dgl/graphbolt/sampled_subgraph.py