def sample(edge_index, num,num_nodes,sizes): edge_index = edge_index.to('cpu') value = torch.arange(edge_index.size(1)) adj_t = SparseTensor(row=edge_index[0], col=edge_index[1], value=value, sparse_sizes=(num_nodes, num_nodes)).t() adjs = [] idx = torch.tensor(num, dtype=torch.long) n_id = idx for size in sizes: adj_t, n_id = adj_t.sample_adj(n_id, size, replace=False) e_id = adj_t.storage.value() size = adj_t.sparse_sizes()[::-1] row, col, _ = adj_t.coo() edge_index = torch.stack([col, row], dim=0) adjs.append(EdgeIndex(edge_index, e_id, size)) adjs = adjs[0] if len(adjs) == 1 else adjs[::-1] return n_id,adjs

这段代码实现了一个采样邻居节点的函数。具体来说，输入参数包括：

- edge_index: 表示图的边的索引，大小为(2, num_edges)，其中第一行表示源节点，第二行表示目标节点。
- num: 表示采样的起始节点索引，大小为1。
- num_nodes: 表示图的节点总数。
- sizes: 表示需要采样的邻居节点数，是一个长度为K的列表。

函数的输出包括：

- n_id: 表示采样后的节点索引，大小为M。
- adjs: 表示采样后得到的邻接矩阵，是一个EdgeIndex对象，包含三个属性：edge_index表示边的索引，edge_attr表示边的属性，size表示邻接矩阵的大小。

在函数内部，首先将edge_index转移到CPU上，并使用torch.arange函数创建了一个大小为num_edges的张量value，表示每条边的编号。然后使用SparseTensor函数将edge_index和value转换为一个稀疏矩阵adj_t，其中sparse_sizes表示邻接矩阵的大小。接着，定义一个空列表adjs，表示存储采样后的邻接矩阵。然后用idx表示当前采样的节点索引，初始化为num。接下来，对于sizes列表中的每一个元素size，使用sample_adj函数对邻接矩阵adj_t进行采样，得到新的邻接矩阵adj_t和节点索引n_id。然后从adj_t中获取采样后的边的索引e_id，邻接矩阵大小size，以及列索引col和行索引row，根据col和row构建新的边的索引edge_index，并将结果存储到列表adjs中。最后，根据adjs的长度返回结果，如果adjs长度为1，则返回第一个元素，否则将adjs列表反转并返回。