Repository: pyg-team/pytorch_geometric
Pull request #8684 (2023-12-29 11:13:10)

`EdgeIndex` support in `Dataset` and `Batch`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8684?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`4b8a866`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/4b8a866bcc72394277baef852c56b3bcf90c3c3f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.47% compared to head [(`2a37231`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8684?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.48%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8684?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/data/separate.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8684?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvc2VwYXJhdGUucHk=) | 95.65% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8684?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8684   +/-   ##
=======================================
  Coverage   89.47%   89.48%           
=======================================
  Files         481      481           
  Lines       30992    31009   +17     
=======================================
+ Hits        27731    27747   +16     
- Misses       3261     3262    +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8684?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/torch_geometric/data/separate.py b/torch_geometric/data/separate.py
index da5f7653c..d08454e23 100644
--- a/torch_geometric/data/separate.py
+++ b/torch_geometric/data/separate.py
@@ -1,16 +1,25 @@
 from collections.abc import Mapping, Sequence
-from typing import Any
+from typing import Any, Type, TypeVar
 
 from torch import Tensor
 
+from torch_geometric import EdgeIndex
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage
 from torch_geometric.typing import SparseTensor, TensorFrame
 from torch_geometric.utils import narrow
 
+T = TypeVar('T')
 
-def separate(cls, batch: Any, idx: int, slice_dict: Any, inc_dict: Any = None,
-             decrement: bool = True) -> BaseData:
+
+def separate(
+    cls: Type[T],
+    batch: Any,
+    idx: int,
+    slice_dict: Any,
+    inc_dict: Any = None,
+    decrement: bool = True,
+) -> T:
     # Separates the individual element from a `batch` at index `idx`.
     # `separate` can handle both homogeneous and heterogeneous data objects by
     # individually separating all their stores.
@@ -19,15 +28,15 @@ def separate(cls, batch: Any, idx: int, slice_dict: Any, inc_dict: Any = None,
 
     data = cls().stores_as(batch)
 
-    # We iterate over each storage object and recursively separate all its
-    # attributes:
+    # Iterate over each storage object and recursively separate its attributes:
     for batch_store, data_store in zip(batch.stores, data.stores):
         key = batch_store._key
-        if key is not None:
+        if key is not None:  # Heterogeneous:
             attrs = slice_dict[key].keys()
-        else:
+        else:  # Homogeneous:
             attrs = set(batch_store.keys())
             attrs = [attr for attr in slice_dict.keys() if attr in attrs]
+
         for attr in attrs:
             if key is not None:
                 slices = slice_dict[key][attr]
@@ -35,6 +44,7 @@ def separate(cls, batch: Any, idx: int, slice_dict: Any, inc_dict: Any = None,
             else:
                 slices = slice_dict[attr]
                 incs = inc_dict[attr] if decrement else None
+
             data_store[attr] = _separate(attr, batch_store[attr], idx, slices,
                                          incs, batch, batch_store, decrement)
 
@@ -48,7 +58,7 @@ def separate(cls, batch: Any, idx: int, slice_dict: Any, inc_dict: Any = None,
 
 def _separate(
     key: str,
-    value: Any,
+    values: Any,
     idx: int,
     slices: Any,
     incs: Any,
@@ -57,43 +67,51 @@ def _separate(
     decrement: bool,
 ) -> Any:
 
-    if isinstance(value, Tensor):
+    if isinstance(values, Tensor):
         # Narrow a `torch.Tensor` based on `slices`.
         # NOTE: We need to take care of decrementing elements appropriately.
         key = str(key)
-        cat_dim = batch.__cat_dim__(key, value, store)
+        cat_dim = batch.__cat_dim__(key, values, store)
         start, end = int(slices[idx]), int(slices[idx + 1])
-        value = narrow(value, cat_dim or 0, start, end - start)
+        value = narrow(values, cat_dim or 0, start, end - start)
         value = value.squeeze(0) if cat_dim is None else value
+
+        if isinstance(values, EdgeIndex) and values._cat_metadata is not None:
+            # Reconstruct original `EdgeIndex` metadata:
+            value._sparse_size = values._cat_metadata.sparse_size[idx]
+            value._sort_order = values._cat_metadata.sort_order[idx]
+            value._is_undirected = values._cat_metadata.is_undirected[idx]
+
         if (decrement and incs is not None
                 and (incs.dim() > 1 or int(incs[idx]) != 0)):
             value = value - incs[idx].to(value.device)
+
         return value
 
-    elif isinstance(value, SparseTensor) and decrement:
+    elif isinstance(values, SparseTensor) and decrement:
         # Narrow a `SparseTensor` based on `slices`.
         # NOTE: `cat_dim` may return a tuple to allow for diagonal stacking.
         key = str(key)
-        cat_dim = batch.__cat_dim__(key, value, store)
+        cat_dim = batch.__cat_dim__(key, values, store)
         cat_dims = (cat_dim, ) if isinstance(cat_dim, int) else cat_dim
         for i, dim in enumerate(cat_dims):
             start, end = int(slices[idx][i]), int(slices[idx + 1][i])
-            value = value.narrow(dim, start, end - start)
-        return value
+            values = values.narrow(dim, start, end - start)
+        return values
 
-    elif isinstance(value, TensorFrame):
+    elif isinstance(values, TensorFrame):
         key = str(key)
         start, end = int(slices[idx]), int(slices[idx + 1])
-        value = value[start:end]
+        value = values[start:end]
         return value
 
-    elif isinstance(value, Mapping):
+    elif isinstance(values, Mapping):
         # Recursively separate elements of dictionaries.
         return {
             key:
             _separate(
                 key,
-                elem,
+                value,
                 idx,
                 slices=slices[key],
                 incs=incs[key] if decrement else None,
@@ -101,32 +119,32 @@ def _separate(
                 store=store,
                 decrement=decrement,
             )
-            for key, elem in value.items()
+            for key, value in values.items()
         }
 
-    elif (isinstance(value, Sequence) and isinstance(value[0], Sequence)
-          and not isinstance(value[0], str) and len(value[0]) > 0
-          and isinstance(value[0][0], (Tensor, SparseTensor))
+    elif (isinstance(values, Sequence) and isinstance(values[0], Sequence)
+          and not isinstance(values[0], str) and len(values[0]) > 0
+          and isinstance(values[0][0], (Tensor, SparseTensor))
           and isinstance(slices, Sequence)):
         # Recursively separate elements of lists of lists.
-        return [elem[idx] for elem in value]
+        return [value[idx] for value in values]
 
-    elif (isinstance(value, Sequence) and not isinstance(value, str)
-          and isinstance(value[0], (Tensor, SparseTensor))
+    elif (isinstance(values, Sequence) and not isinstance(values, str)
+          and isinstance(values[0], (Tensor, SparseTensor))
           and isinstance(slices, Sequence)):
         # Recursively separate elements of lists of Tensors/SparseTensors.
         return [
             _separate(
                 key,
-                elem,
+                value,
                 idx,
                 slices=slices[i],
                 incs=incs[i] if decrement else None,
                 batch=batch,
                 store=store,
                 decrement=decrement,
-            ) for i, elem in enumerate(value)
+            ) for i, value in enumerate(values)
         ]
 
     else:
-        return value[idx]
+        return values[idx]
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index 0fab30d29..4d3b04dbc 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -54,6 +54,13 @@ class SortOrder(Enum):
     COL = 'col'
 
 
+class CatMetadata(NamedTuple):
+    nnz: List[int]
+    sparse_size: List[Tuple[Optional[int], Optional[int]]]
+    sort_order: List[Optional[SortOrder]]
+    is_undirected: List[bool]
+
+
 def implements(torch_function: Callable) -> Callable:
     r"""Registers a :pytorch:`PyTorch` function override."""
     @functools.wraps(torch_function)
@@ -239,6 +246,10 @@ class EdgeIndex(Tensor):
     # A cached "1"-value vector for `torch.sparse` matrix multiplication:
     _value: Optional[Tensor] = None
 
+    # Whenever we perform a concatenation of edge indices, we cache the
+    # original metadata to be able to reconstruct individual edge indices:
+    _cat_metadata: Optional[CatMetadata] = None
+
     def __new__(
         cls: Type,
         data: Any,
@@ -1030,30 +1041,39 @@ def cat(
 
     output = output.as_subclass(EdgeIndex)
 
+    nnz_list = [t.size(1) for t in tensors]
+    sparse_size_list = [t.sparse_size() for t in tensors]  # type: ignore
+    sort_order_list = [t._sort_order for t in tensors]  # type: ignore
+    is_undirected_list = [t.is_undirected for t in tensors]  # type: ignore
+
     # Post-process `sparse_size`:
-    num_rows: Optional[int] = 0
-    for tensor in tensors:
-        if not isinstance(tensor, EdgeIndex) or tensor.num_rows is None:
-            num_rows = None
+    total_num_rows: Optional[int] = 0
+    for num_rows, _ in sparse_size_list:
+        if num_rows is None:
+            total_num_rows = None
             break
-        assert isinstance(num_rows, int)
-        num_rows = max(num_rows, tensor.num_rows)
+        assert isinstance(total_num_rows, int)
+        total_num_rows = max(num_rows, total_num_rows)
 
-    num_cols: Optional[int] = 0
-    for tensor in tensors:
-        if not isinstance(tensor, EdgeIndex) or tensor.num_cols is None:
-            num_cols = None
+    total_num_cols: Optional[int] = 0
+    for _, num_cols in sparse_size_list:
+        if num_cols is None:
+            total_num_cols = None
             break
-        assert isinstance(num_cols, int)
-        num_cols = max(num_cols, tensor.num_cols)
+        assert isinstance(total_num_cols, int)
+        num_cols = max(num_cols, total_num_cols)
 
     output._sparse_size = (num_rows, num_cols)
 
     # Post-process `is_undirected`:
-    output._is_undirected = all([
-        isinstance(tensor, EdgeIndex) and tensor.is_undirected
-        for tensor in tensors
-    ])
+    output._is_undirected = all(is_undirected_list)
+
+    output._cat_metadata = CatMetadata(
+        nnz=nnz_list,
+        sparse_size=sparse_size_list,
+        sort_order=sort_order_list,
+        is_undirected=is_undirected_list,
+    )
 
     return output
 


Test patch:
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
index 5ab69df63..8d66bff81 100644
--- a/test/data/test_batch.py
+++ b/test/data/test_batch.py
@@ -5,6 +5,7 @@ import pytest
 import torch
 
 import torch_geometric
+from torch_geometric import EdgeIndex
 from torch_geometric.data import Batch, Data, HeteroData
 from torch_geometric.testing import get_random_edge_index, withPackage
 from torch_geometric.typing import SparseTensor
@@ -107,6 +108,46 @@ def test_batch_basic():
     torch_geometric.set_debug(True)
 
 
+def test_edge_index():
+    edge_index1 = EdgeIndex(
+        [[0, 1, 1, 2], [1, 0, 2, 1]],
+        sparse_size=(3, 3),
+        sort_order='row',
+        is_undirected=True,
+    )
+    edge_index2 = EdgeIndex(
+        [[1, 0, 2, 1, 3, 2], [0, 1, 1, 2, 2, 3]],
+        sparse_size=(4, 4),
+        sort_order='col',
+    )
+
+    data1 = Data(edge_index=edge_index1)
+    data2 = Data(edge_index=edge_index2)
+
+    batch = Batch.from_data_list([data1, data2])
+
+    assert len(batch) == 2
+    assert batch.batch.equal(torch.tensor([0, 0, 0, 1, 1, 1, 1]))
+    assert batch.ptr.equal(torch.tensor([0, 3, 7]))
+    assert isinstance(batch.edge_index, EdgeIndex)
+    assert batch.edge_index.equal(
+        torch.tensor([
+            [0, 1, 1, 2, 4, 3, 5, 4, 6, 5],
+            [1, 0, 2, 1, 3, 4, 4, 5, 5, 6],
+        ]))
+    assert batch.edge_index.sparse_size() == (7, 7)
+    assert batch.edge_index.sort_order is None
+    assert not batch.edge_index.is_undirected
+
+    for i, edge_index in enumerate([edge_index1, edge_index2]):
+        data = batch[i]
+        assert isinstance(data.edge_index, EdgeIndex)
+        assert data.edge_index.equal(edge_index)
+        assert data.edge_index.sparse_size() == edge_index.sparse_size()
+        assert data.edge_index.sort_order == edge_index.sort_order
+        assert data.edge_index.is_undirected == edge_index.is_undirected
+
+
 @withPackage('torch_sparse')
 def test_batch_with_sparse_tensor():
     x = SparseTensor.from_dense(torch.tensor([[1.0], [2.0], [3.0]]))
diff --git a/test/data/test_dataset.py b/test/data/test_dataset.py
index c6ad968dd..a7bbe048f 100644
--- a/test/data/test_dataset.py
+++ b/test/data/test_dataset.py
@@ -3,6 +3,7 @@ import copy
 import pytest
 import torch
 
+from torch_geometric import EdgeIndex
 from torch_geometric.data import Data, HeteroData, InMemoryDataset
 from torch_geometric.datasets import KarateClub
 from torch_geometric.testing import withPackage
@@ -121,6 +122,41 @@ def test_stored_hetero_in_memory_dataset(tmp_path):
     assert dataset[1]['paper'].num_nodes == 4
 
 
+def test_edge_index(tmp_path):
+    edge_index1 = EdgeIndex(
+        [[0, 1, 1, 2], [1, 0, 2, 1]],
+        sparse_size=(3, 3),
+        sort_order='row',
+        is_undirected=True,
+    )
+    edge_index2 = EdgeIndex(
+        [[1, 0, 2, 1, 3, 2], [0, 1, 1, 2, 2, 3]],
+        sparse_size=(4, 4),
+        sort_order='col',
+    )
+
+    data1 = Data(edge_index=edge_index1)
+    data2 = Data(edge_index=edge_index2)
+
+    dataset = MyTestDataset([data1, data2])
+    assert len(dataset) == 2
+    for data, edge_index in zip(dataset, [edge_index1, edge_index2]):
+        assert isinstance(data.edge_index, EdgeIndex)
+        assert data.edge_index.equal(edge_index)
+        assert data.edge_index.sparse_size() == edge_index.sparse_size()
+        assert data.edge_index.sort_order == edge_index.sort_order
+        assert data.edge_index.is_undirected == edge_index.is_undirected
+
+    dataset = MyStoredTestDataset(tmp_path, [data1, data2])
+    assert len(dataset) == 2
+    for data, edge_index in zip(dataset, [edge_index1, edge_index2]):
+        assert isinstance(data.edge_index, EdgeIndex)
+        assert data.edge_index.equal(edge_index)
+        assert data.edge_index.sparse_size() == edge_index.sparse_size()
+        assert data.edge_index.sort_order == edge_index.sort_order
+        assert data.edge_index.is_undirected == edge_index.is_undirected
+
+
 def test_in_memory_num_classes():
     dataset = MyTestDataset([Data(), Data()])
     assert dataset.num_classes == 0
@@ -129,7 +165,8 @@ def test_in_memory_num_classes():
     assert dataset.num_classes == 2
 
     dataset = MyTestDataset([Data(y=1.5), Data(y=2.5), Data(y=3.5)])
-    assert dataset.num_classes == 3
+    with pytest.warns(UserWarning, match="unique elements"):
+        assert dataset.num_classes == 3
 
     dataset = MyTestDataset([
         Data(y=torch.tensor([[0, 1, 0, 1]])),
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 213f910c7..f8b567320 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -371,6 +371,11 @@ def test_cat(dtype, device, is_undirected):
     assert not out.is_sorted
     assert out.is_undirected == is_undirected
 
+    assert out._cat_metadata.nnz == [4, 4]
+    assert out._cat_metadata.sparse_size == [(3, 3), (4, 4)]
+    assert out._cat_metadata.sort_order == [None, None]
+    assert out._cat_metadata.is_undirected == [is_undirected, is_undirected]
+
     out = torch.cat([adj1, adj2, adj3], dim=1)
     assert out.size() == (2, 12)
     assert isinstance(out, EdgeIndex)