Add unbatch functionality

torch_geometric/loader/utils.py

@@ -2,6 +2,7 @@
 import math
 from typing import Dict, Optional, Tuple, Union
 
+import pandas
 import torch
 from torch import Tensor
 from torch_sparse import SparseTensor
@@ -24,6 +25,22 @@ def index_select(value: Tensor, index: Tensor, dim: int = 0) -> Tensor:
     return torch.index_select(value, 0, index, out=out)
 
 
+def unbatcher(src: torch.tensor, index: torch.tensor, dim: int = -1) -> list:
+    scatter = pandas.DataFrame(
+        index, columns=['scatter']).reset_index().groupby('scatter').tail(1)
+    indices = []
+    index_i = 0  # index starts for a graph
+    index_f = 0  # index ends for a graph
+    for i in range(0, len(scatter)):
+        index_f = scatter['index'].iloc[i] + 1
+        indices.append(torch.tensor(range(index_i, index_f)))
+        index_i = index_f
+    unbatched = []
+    for i in range(0, len(indices)):
+        unbatched.append(torch.index_select(src, dim, indices[i]))
+    return unbatched
+
+
 def edge_type_to_str(edge_type: Union[EdgeType, str]) -> str:
     # Since C++ cannot take dictionaries with tuples as key as input, edge type
     # triplets need to be converted into single strings.