unbatch_edge_index drops trailing edgeless graphs when batch_size is inferred

Author: HaokaiDingCreated Sep 16, 2026Updated Sep 16, 2026

Describe the bug

When batch_size is omitted, unbatch_edge_index() drops trailing graphs that have nodes but no edges. These graphs are present in the node batch vector, so their empty edge partitions can be inferred. If every graph has no edges, the function returns no partitions at all.

python
import torch
from torch_geometric.data import Batch, Data
from torch_geometric.utils import unbatch_edge_index

batch = Batch.from_data_list([
    Data(edge_index=torch.tensor([[0, 1], [1, 0]]), num_nodes=2),
    Data(edge_index=torch.empty((2, 0), dtype=torch.long), num_nodes=1),
])
edge_indices = unbatch_edge_index(batch.edge_index, batch.batch)
print(batch.num_graphs, len(edge_indices))  # 2 1
assert len(edge_indices) == batch.num_graphs

Expected: two partitions, with the second tensor having shape (2, 0). Actual: only the first graph's edges are returned and the assertion fails.

Passing batch_size=batch.num_graphs is a workaround. Internally, the first degree(batch, ...) already determines the correct graph count, but the second call infers it again from edge_batch, which cannot represent trailing edgeless graphs. Reusing the first degree vector's length fixes the mismatch.

This report concerns graphs represented in the node batch vector. Trailing graphs with zero nodes still need an explicit batch_size.

Versions

  • Current master source, reporting PyG 2.9.0.
  • PyTorch 2.14.0, Python 3.13.15, NumPy 2.5.3.
  • macOS 26.6.2 ARM64; reproduced on CPU with synthetic graphs and no external data.

Source: pyg-team/pytorch_geometric