Repository: pyg-team/pytorch_geometric
Pull request #8871 (2024-02-06 11:57:33)

Make `Data` attributes cache non shareable
When `Data` object has non-empty `_attr_cache`, after creating copy of `Data`, both objects will share `_attr_cache`, which may result in wrong attribute classification. For example:
```python
x = torch.rand((4, 4))
edge_index = torch.tensor([[0, 1, 2, 3, 0, 1], [0, 1, 2, 3, 0, 1]], dtype=torch.long)
time = torch.tensor([0., 1., 2., 3., 4., 5.])
data = Data(x=x, edge_index=edge_index, time=time)

# this will implicitly fill _attr_cache
assert data.is_node_attr('x')

out = data.up_to(3.5)
assert out.is_node_attr('time')
# before this PR, assert will fail, producing the wrong attribute classification, as time should be an edge attribute
assert not data.is_node_attr('time')
```

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8871?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`64a1268`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/64a12680315dbb5bcd00c4e09cf836b88373d3b3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.28% compared to head [(`4d97f89`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8871?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.28%.


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


```diff
@@           Coverage Diff           @@
##           master    #8871   +/-   ##
=======================================
  Coverage   89.28%   89.28%           
=======================================
  Files         467      467           
  Lines       29900    29901    +1     
=======================================
+ Hits        26696    26697    +1     
  Misses       3204     3204           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8871?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/storage.py b/torch_geometric/data/storage.py
index 1a815a3a7..89a10053b 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -135,7 +135,8 @@ class BaseStorage(MutableMapping):
     def __copy__(self) -> Self:
         out = self.__class__.__new__(self.__class__)
         for key, value in self.__dict__.items():
-            out.__dict__[key] = value
+            if key != '_cached_attr':
+                out.__dict__[key] = value
         out._mapping = copy.copy(out._mapping)
         return out
 
diff --git a/torch_geometric/loader/node_loader.py b/torch_geometric/loader/node_loader.py
index 2c040ae1f..c26c8ccc3 100644
--- a/torch_geometric/loader/node_loader.py
+++ b/torch_geometric/loader/node_loader.py
@@ -196,6 +196,12 @@ class NodeLoader(
             data.num_sampled_nodes = out.num_sampled_nodes
             data.num_sampled_edges = out.num_sampled_edges
 
+            if out.orig_row is not None and out.orig_col is not None:
+                data._orig_edge_index = torch.stack([
+                    out.orig_row,
+                    out.orig_col,
+                ], dim=0)
+
             data.input_id = out.metadata[0]
             data.seed_time = out.metadata[1]
             data.batch_size = out.metadata[0].size(0)
@@ -238,6 +244,13 @@ class NodeLoader(
             data.set_value_dict('num_sampled_nodes', out.num_sampled_nodes)
             data.set_value_dict('num_sampled_edges', out.num_sampled_edges)
 
+            if out.orig_row is not None and out.orig_col is not None:
+                for key in out.orig_row.keys():
+                    data[key]._orig_edge_index = torch.stack([
+                        out.orig_row[key],
+                        out.orig_col[key],
+                    ], dim=0)
+
             input_type = self.input_data.input_type
             data[input_type].input_id = out.metadata[0]
             data[input_type].seed_time = out.metadata[1]
diff --git a/torch_geometric/sampler/base.py b/torch_geometric/sampler/base.py
index 99f1b4a5a..6332b4073 100644
--- a/torch_geometric/sampler/base.py
+++ b/torch_geometric/sampler/base.py
@@ -183,6 +183,14 @@ class SamplerOutput(CastMixin):
             per hop. (default: :obj:`None`)
         num_sampled_edges (List[int], optional): The number of sampled edges
             per hop. (default: :obj:`None`)
+        orig_row (torch.Tensor, optional): The original source node indices
+            returned by the sampler.
+            Filled in case :meth:`to_bidirectional` is called with the
+            :obj:`keep_orig_edges` option. (default: :obj:`None`)
+        orig_col (torch.Tensor, optional): The original destination node
+            indices indices returned by the sampler.
+            Filled in case :meth:`to_bidirectional` is called with the
+            :obj:`keep_orig_edges` option. (default: :obj:`None`)
         metadata: (Any, optional): Additional metadata information.
             (default: :obj:`None`)
     """
@@ -193,16 +201,31 @@ class SamplerOutput(CastMixin):
     batch: OptTensor = None
     num_sampled_nodes: Optional[List[int]] = None
     num_sampled_edges: Optional[List[int]] = None
+    orig_row: Tensor = None
+    orig_col: Tensor = None
     # TODO(manan): refine this further; it does not currently define a proper
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
 
-    def to_bidirectional(self) -> 'SamplerOutput':
+    def to_bidirectional(
+        self,
+        keep_orig_edges: bool = False,
+    ) -> 'SamplerOutput':
         r"""Converts the sampled subgraph into a bidirectional variant, in
         which all sampled edges are guaranteed to be bidirectional.
+
+        Args:
+            keep_orig_edges (bool, optional): If specified, directional edges
+                are still maintained. (default: :obj:`False`)
         """
         out = copy.copy(self)
 
+        if keep_orig_edges:
+            out.orig_row = self.row
+            out.orig_col = self.col
+        else:
+            out.num_sampled_nodes = out.num_sampled_edges = None
+
         out.row, out.col, out.edge = to_bidirectional(
             row=self.row,
             col=self.col,
@@ -211,7 +234,6 @@ class SamplerOutput(CastMixin):
             edge_id=self.edge,
             rev_edge_id=self.edge,
         )
-        out.num_sampled_nodes = out.num_sampled_edges = None
 
         return out
 
@@ -248,6 +270,14 @@ class HeteroSamplerOutput(CastMixin):
         num_sampled_edges (Dict[EdgeType, List[int]], optional): The number of
             sampled edges for each edge type and each layer.
             (default: :obj:`None`)
+        orig_row (Dict[EdgeType, torch.Tensor], optional): The original source
+            node indices returned by the sampler.
+            Filled in case :meth:`to_bidirectional` is called with the
+            :obj:`keep_orig_edges` option. (default: :obj:`None`)
+        orig_col (Dict[EdgeType, torch.Tensor], optional): The original
+            destination node indices returned by the sampler.
+            Filled in case :meth:`to_bidirectional` is called with the
+            :obj:`keep_orig_edges` option. (default: :obj:`None`)
         metadata: (Any, optional): Additional metadata information.
             (default: :obj:`None`)
     """
@@ -258,19 +288,37 @@ class HeteroSamplerOutput(CastMixin):
     batch: Optional[Dict[NodeType, Tensor]] = None
     num_sampled_nodes: Optional[Dict[NodeType, List[int]]] = None
     num_sampled_edges: Optional[Dict[EdgeType, List[int]]] = None
+    orig_row: Optional[Dict[EdgeType, Tensor]] = None
+    orig_col: Optional[Dict[EdgeType, Tensor]] = None
     # TODO(manan): refine this further; it does not currently define a proper
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
 
-    def to_bidirectional(self) -> 'SamplerOutput':
+    def to_bidirectional(
+        self,
+        keep_orig_edges: bool = False,
+    ) -> 'SamplerOutput':
         r"""Converts the sampled subgraph into a bidirectional variant, in
         which all sampled edges are guaranteed to be bidirectional.
+
+        Args:
+            keep_orig_edges (bool, optional): If specified, directional edges
+                are still maintained. (default: :obj:`False`)
         """
         out = copy.copy(self)
         out.row = copy.copy(self.row)
         out.col = copy.copy(self.col)
         out.edge = copy.copy(self.edge)
 
+        if keep_orig_edges:
+            out.orig_row = {}
+            out.orig_col = {}
+            for key in self.row.keys():
+                out.orig_row[key] = self.row[key]
+                out.orig_col[key] = self.col[key]
+        else:
+            out.num_sampled_nodes = out.num_sampled_edges = None
+
         src_dst_dict = defaultdict(list)
         edge_types = self.row.keys()
         edge_types = [k for k in edge_types if not k[1].startswith('rev_')]
@@ -325,8 +373,6 @@ class HeteroSamplerOutput(CastMixin):
                                   f"since the edge type {edge_type} does not "
                                   f"seem to have a reverse edge type")
 
-        out.num_sampled_nodes = out.num_sampled_edges = None
-
         return out
 
 


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index a895c916f..3b94d4169 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -177,6 +177,20 @@ def test_data_attr_cache():
     assert 'y' in data._store._cached_attr[AttrType.OTHER]
 
 
+def test_data_attr_cache_not_shared():
+    x = torch.rand((4, 4))
+    edge_index = torch.tensor([[0, 1, 2, 3, 0, 1], [0, 1, 2, 3, 0, 1]])
+    time = torch.arange(edge_index.size(1))
+    data = Data(x=x, edge_index=edge_index, time=time)
+    assert data.is_node_attr('x')
+
+    out = data.up_to(3.5)
+    # This is expected behavior due to the ambiguity of between node-level and
+    # edge-level tensors when they share the same number of nodes/edges.
+    assert out.is_node_attr('time')
+    assert not data.is_node_attr('time')
+
+
 def test_to_heterogeneous_empty_edge_index():
     data = Data(
         x=torch.randn(5, 10),