Repository: pyg-team/pytorch_geometric
Pull request #9284 (2024-05-04 16:35:05)

Add `torch.cat` support for `torch_geometric.Index`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 871c3bba5..0921d820f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added `torch_geometric.Index` ([#9276](https://github.com/pyg-team/pytorch_geometric/pull/9276), [#9277](https://github.com/pyg-team/pytorch_geometric/pull/9277), [#9278](https://github.com/pyg-team/pytorch_geometric/pull/9278), [#9279](https://github.com/pyg-team/pytorch_geometric/pull/9279), [#9280](https://github.com/pyg-team/pytorch_geometric/pull/9280), [#9281](https://github.com/pyg-team/pytorch_geometric/pull/9281))
+- Added `torch_geometric.Index` ([#9276](https://github.com/pyg-team/pytorch_geometric/pull/9276), [#9277](https://github.com/pyg-team/pytorch_geometric/pull/9277), [#9278](https://github.com/pyg-team/pytorch_geometric/pull/9278), [#9279](https://github.com/pyg-team/pytorch_geometric/pull/9279), [#9280](https://github.com/pyg-team/pytorch_geometric/pull/9280), [#9281](https://github.com/pyg-team/pytorch_geometric/pull/9281), [#9284](https://github.com/pyg-team/pytorch_geometric/pull/9284))
 - Added support for PyTorch 2.3 ([#9240](https://github.com/pyg-team/pytorch_geometric/pull/9240))
 - Added support for `EdgeIndex` in `message_and_aggregate` ([#9131](https://github.com/pyg-team/pytorch_geometric/pull/9131))
 - Added `CornellTemporalHyperGraphDataset` ([#9090](https://github.com/pyg-team/pytorch_geometric/pull/9090))
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index f7926443d..6c8afc9cc 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1358,9 +1358,9 @@ def _cat(
             total_num_cols = None
             break
         assert isinstance(total_num_cols, int)
-        num_cols = max(num_cols, total_num_cols)
+        total_num_cols = max(num_cols, total_num_cols)
 
-    out._sparse_size = (num_rows, num_cols)
+    out._sparse_size = (total_num_rows, total_num_cols)
 
     # Post-process `is_undirected`:
     out._is_undirected = all(is_undirected_list)
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index aff331a4c..42dc69244 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -459,3 +459,41 @@ def _sort_stable(
         out._is_sorted = True
 
     return out, perm
+
+
+@implements(aten.cat.default)
+def _cat(
+    tensors: List[Union[Index, Tensor]],
+    dim: int = 0,
+) -> Union[Index, Tensor]:
+
+    data_list = pytree.tree_map_only(Index, lambda x: x._data, tensors)
+    data = aten.cat.default(data_list, dim=dim)
+
+    if any([not isinstance(tensor, Index) for tensor in tensors]):
+        return data
+
+    out = Index(data)
+
+    nnz_list = [t.numel() for t in tensors]
+    dim_size_list = [t.dim_size for t in tensors]  # type: ignore
+    is_sorted_list = [t.is_sorted for t in tensors]  # type: ignore
+
+    # Post-process `dim_size`:
+    total_dim_size: Optional[int] = 0
+    for dim_size in dim_size_list:
+        if dim_size is None:
+            total_dim_size = None
+            break
+        assert isinstance(total_dim_size, int)
+        total_dim_size = max(dim_size, total_dim_size)
+
+    out._dim_size = total_dim_size
+
+    out._cat_metadata = CatMetadata(
+        nnz=nnz_list,
+        dim_size=dim_size_list,
+        is_sorted=is_sorted_list,
+    )
+
+    return out


Test patch:
diff --git a/test/test_index.py b/test/test_index.py
index d57eed322..3e930fcc2 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -268,3 +268,40 @@ def test_sort_stable(dtype, device):
     assert perm.equal(tensor([3, 1, 2, 0], device=device))
     assert index.dim_size == 3
     assert not index.is_sorted
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_cat(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index1 = Index([0, 1, 1, 2], dim_size=3, is_sorted=True, **kwargs)
+    index2 = Index([1, 2, 2, 3], dim_size=4, is_sorted=True, **kwargs)
+    index3 = Index([1, 2, 2, 3], **kwargs)
+
+    out = torch.cat([index1, index2])
+    assert out.equal(tensor([0, 1, 1, 2, 1, 2, 2, 3], device=device))
+    assert out.size() == (8, )
+    assert isinstance(out, Index)
+    assert out.dim_size == 4
+    assert not out.is_sorted
+
+    assert out._cat_metadata.nnz == [4, 4]
+    assert out._cat_metadata.dim_size == [3, 4]
+    assert out._cat_metadata.is_sorted == [True, True]
+
+    out = torch.cat([index1, index2, index3])
+    assert out.size() == (12, )
+    assert isinstance(out, Index)
+    assert out.dim_size is None
+    assert not out.is_sorted
+
+    out = torch.cat([index1, index2.as_tensor()])
+    assert out.size() == (8, )
+    assert not isinstance(out, Index)
+
+    inplace = torch.empty(8, dtype=dtype, device=device)
+    out = torch.cat([index1, index2], out=inplace)
+    assert out.equal(tensor([0, 1, 1, 2, 1, 2, 2, 3], device=device))
+    assert out.data_ptr() == inplace.data_ptr()
+    assert not isinstance(out, Index)
+    assert not isinstance(inplace, Index)