unbatch_edge_index returns less splits than number of batches

[hcagri]

🐛 Describe the bug

The issue I am encountering with the unbatch_edge_index function from PyTorch Geometric is that there is a discrepancy between the number of batches I provide and the number of edge indices returned by the function.

Currently, the function seems to omit empty edge_index tensors for the last batch if it does not have any edges. As a result, the number of returned edge indices is less than the total number of batches.

The problem can be reproduced by the simple script provided below.

from torch_geometric.utils import  unbatch_edge_index
import torch 

x = torch.tensor([0,1,2,3,4,5])
batch = torch.tensor([0, 0, 1, 1, 2, 2])
edge_index = torch.tensor([[0,2],[1,3]])
edge_index_splits = unbatch_edge_index(edge_index, batch)

print("Number of Batches: ", len(torch.unique(batch)))
print("Number of Edge Indices: ", len(edge_index_splits))

Environment

• PyG version: 2.3.0
• PyTorch version: 2.0.1
• OS: Ubuntu 20.04.2 LTS
• Python Version: 3.9.16

[rusty1s]

Thanks for reporting. Will be fixed via #7851.