Add unbatch functionality

test/utils/test_unbatch.py

@@ -0,0 +1,13 @@
+import torch
+
+from torch_geometric.utils import unbatch
+
+
+def test_unbatch():
+    src = torch.arange(10)
+    batch = torch.tensor([0, 0, 0, 1, 1, 2, 2, 3, 4, 4])
+
+    out = unbatch(src, batch)
+    assert len(out) == 5
+    for i in range(len(out)):
+        assert torch.equal(out[i], src[batch == i])

torch_geometric/utils/__init__.py

@@ -17,6 +17,7 @@
 from .to_dense_batch import to_dense_batch
 from .to_dense_adj import to_dense_adj
 from .sparse import dense_to_sparse
+from .unbatch import unbatch
 from .normalized_cut import normalized_cut
 from .grid import grid
 from .geodesic import geodesic_distance
@@ -61,6 +62,7 @@
     'to_dense_batch',
     'to_dense_adj',
     'dense_to_sparse',
+    'unbatch',
     'normalized_cut',
     'grid',
     'geodesic_distance',

torch_geometric/utils/unbatch.py

@@ -0,0 +1,23 @@
+from typing import List
+
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import degree
+
+
+def unbatch(src: Tensor, batch: Tensor, dim: int = 0) -> List[Tensor]:
+    r"""Splits :obj:`src` according to a :obj:`batch` vector along dimension
+    :obj:`dim`.
+
+    Args:
+        src (Tensor): The source tensor.
+        batch (LongTensor): The batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            entry in :obj:`src` to a specific example. Must be ordered.
+        dim (int): The dimension along which to split the :obj:`src` tensor.
+
+    :rtype: :class:`List[Tensor]`
+    """
+    sizes = degree(batch, dtype=torch.long).tolist()
+    return src.split(sizes, dim)