Repository: pyg-team/pytorch_geometric
Pull request #7643 (2023-06-26 13:38:15)

`torch.nested_tensor` support in `Data` and `Batch`
Fixes https://github.com/pyg-team/pytorch_geometric/discussions/7622


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4b51ec309..48f0052c8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `torch.nested_tensor` support in `Data` and `Batch` ([#7643](https://github.com/pyg-team/pytorch_geometric/pull/7643))
 - Added `interval` argument to `Cartesian` and `LocalCartesian` transformations ([#7533](https://github.com/pyg-team/pytorch_geometric/pull/7533), [#7614](https://github.com/pyg-team/pytorch_geometric/pull/7614))
 - Added a `LightGCN` example on the `AmazonBook` dataset ([7603](https://github.com/pyg-team/pytorch_geometric/pull/7603))
 - Added a tutorial on hierarchical neighborhood sampling ([#7594](https://github.com/pyg-team/pytorch_geometric/pull/7594))
diff --git a/torch_geometric/data/collate.py b/torch_geometric/data/collate.py
index a9896916d..cc5c1c644 100644
--- a/torch_geometric/data/collate.py
+++ b/torch_geometric/data/collate.py
@@ -159,7 +159,14 @@ def _collate(
         else:
             out = None
 
-        value = torch.cat(values, dim=cat_dim or 0, out=out)
+        if elem.is_nested:
+            tensors = []
+            for nested_tensor in values:
+                tensors.extend(nested_tensor.unbind())
+            value = torch.nested.nested_tensor(tensors)
+        else:
+            value = torch.cat(values, dim=cat_dim or 0, out=out)
+
         return value, slices, incs
 
     elif is_sparse(elem) and increment:
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 0ed967dd1..a9fa1f237 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -955,6 +955,8 @@ def size_repr(key: Any, value: Any, indent: int = 0) -> str:
     pad = ' ' * indent
     if isinstance(value, Tensor) and value.dim() == 0:
         out = value.item()
+    elif isinstance(value, Tensor) and value.is_nested:
+        out = str(list(value.to_padded_tensor(padding=0.0).size()))
     elif isinstance(value, Tensor):
         out = str(list(value.size()))
     elif isinstance(value, np.ndarray):
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index f43cf1254..b46b0484a 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -296,11 +296,17 @@ class NodeStorage(BaseStorage):
         if 'num_nodes' in self:
             return self['num_nodes']
         for key, value in self.items():
-            if isinstance(value, (Tensor, np.ndarray)) and key in N_KEYS:
+            if isinstance(value, Tensor) and key in N_KEYS:
+                cat_dim = self._parent().__cat_dim__(key, value, self)
+                return value.size(cat_dim)
+            if isinstance(value, np.ndarray) and key in N_KEYS:
                 cat_dim = self._parent().__cat_dim__(key, value, self)
                 return value.shape[cat_dim]
         for key, value in self.items():
-            if isinstance(value, (Tensor, np.ndarray)) and 'node' in key:
+            if isinstance(value, Tensor) and 'node' in key:
+                cat_dim = self._parent().__cat_dim__(key, value, self)
+                return value.size(cat_dim)
+            if isinstance(value, np.ndarray) and 'node' in key:
                 cat_dim = self._parent().__cat_dim__(key, value, self)
                 return value.shape[cat_dim]
         if 'adj' in self and isinstance(self.adj, SparseTensor):


Test patch:
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
index ca5f46324..86f1d523b 100644
--- a/test/data/test_batch.py
+++ b/test/data/test_batch.py
@@ -3,6 +3,7 @@ import os.path as osp
 import numpy as np
 import pytest
 import torch
+from torch.nested import nested_tensor
 
 import torch_geometric
 from torch_geometric.data import Batch, Data, HeteroData
@@ -517,3 +518,19 @@ def test_torch_sparse_batch(layout):
     out = to_edge_index(batch.adj.to_sparse(layout=torch.sparse_csr))
     assert torch.equal(out[0], torch.cat([edge_index, edge_index + 3], 1))
     assert torch.equal(out[1], torch.cat([edge_attr, edge_attr], 0))
+
+
+def test_torch_nested_batch():
+    x1 = nested_tensor([torch.randn(3), torch.randn(4)])
+    data1 = Data(x=x1)
+    assert str(data1) == 'Data(x=[2, 4])'
+
+    x2 = nested_tensor([torch.randn(3), torch.randn(4), torch.randn(5)])
+    data2 = Data(x=x2)
+    assert str(data2) == 'Data(x=[3, 5])'
+
+    batch = Batch.from_data_list([data1, data2])
+    assert str(batch) == 'DataBatch(x=[5, 5], batch=[5], ptr=[3])'
+
+    x = nested_tensor(list(x1.unbind() + x2.unbind())).to_padded_tensor(0.0)
+    assert torch.equal(batch.x.to_padded_tensor(0.0), x)