How to control number of samples per graph in batched_negative_sampling?

[vietngth]

Hi, I would like to perform negative sampling for a batch of undirected graphs. I did the following code:

edge_index1 = torch.as_tensor([[0, 0, 1, 2, 3], [0, 1, 0, 3, 2]])
edge_index2 = torch.as_tensor([[4, 6, 5], [4, 5, 6]])
edge_index = torch.cat([edge_index1, edge_index2], dim=1) # shape [2, 8]
batch = torch.as_tensor([0, 0, 0, 0, 1, 1, 1])
batched_negative_sampling(edge_index, batch, num_neg_samples=edge_index.shape[1]*2, force_undirected=True) # shape [2, 12]

Results:

tensor([[0, 0, 1, 1, 2, 3, 2, 3, 4, 4, 5, 6],
        [2, 3, 2, 3, 0, 0, 1, 1, 5, 6, 4, 4]])

1. It doesn't seem to generate the exact number of negative samples I provided. Could you help me take a look at it?
2. Is there already a support function to generate negative samples as self-loops only?

Thanks!

[rusty1s]

1. I think this works as expected, right? You can only sample 4 valid undirected edges as negatives in graph 1, and you can only sample 2 valid undirected edges as negatives in graph 2, leading to an output shape of [2, 12]. Note that negative sampling is done without replacement, and so num_neg_samples only acts as an upper bound.
2. Currently, negative sampling does not include self-loops. We would need to add support for it.