update

test/loader/test_neighbor_loader.py

@@ -392,20 +392,15 @@ def test_temporal_custom_neighbor_loader_on_cora(get_dataset, FeatureStore,
 
     loader1 = NeighborLoader(hetero_data, num_neighbors=[-1, -1],
                              input_nodes='paper', time_attr='time',
-                             batch_size=1)
+                             batch_size=128)
 
     loader2 = NeighborLoader(
         (feature_store, graph_store),
         num_neighbors=[-1, -1],
         input_nodes=TensorAttr(group_name='paper', attr_name='x'),
         time_attr='time',
-        batch_size=1,
+        batch_size=128,
     )
 
-    num_iter = 10
-    for i, (batch1, batch2) in enumerate(zip(loader1, loader2)):
-        mask1 = batch1['paper'].time[0] >= batch1['paper'].time[1:]
-        mask2 = batch2['paper'].time[0] >= batch2['paper'].time[1:]
-        assert torch.all(mask1 & mask2)
-        if i > num_iter:
-            break
+    for batch1, batch2 in zip(loader1, loader2):
+        assert torch.equal(batch1['paper'].time, batch2['paper'].time)

torch_geometric/data/data.py

@@ -874,10 +874,23 @@ def get_all_edge_attrs(self) -> List[EdgeAttr]:
         in `Data` and their layouts"""
         if not hasattr(self, '_edge_attrs'):
             return []
+        added_attrs = set()
 
+        # Check edges added via _put_edge_index:
         edge_attrs = self._edge_attrs.values()
         for attr in edge_attrs:
             attr.size = (self.num_nodes, self.num_nodes)
+            added_attrs.add(attr.layout)
+
+        # Check edges added through regular interface:
+        # TODO deprecate this and store edge attributes for all edges in
+        # EdgeStorage
+        for layout, attr_name in EDGE_LAYOUT_TO_ATTR_NAME.items():
+            if attr_name in self and layout not in added_attrs:
+                edge_attrs.append(
+                    EdgeAttr(edge_type=None, layout=layout,
+                             size=(self.num_nodes, self.num_nodes)))
+
         return edge_attrs
 
 

torch_geometric/data/hetero_data.py

@@ -787,13 +787,30 @@ def get_all_edge_attrs(self) -> List[EdgeAttr]:
         r"""Returns a list of `EdgeAttr` objects corresponding to the edge
         indices stored in `HeteroData` and their layouts."""
         out = []
+        added_attrs = set()
+
+        # Check edges added via _put_edge_index:
         for edge_type, _ in self.edge_items():
             if not hasattr(self[edge_type], '_edge_attrs'):
                 continue
             edge_attrs = self[edge_type]._edge_attrs.values()
             for attr in edge_attrs:
                 attr.size = self[edge_type].size()
+                added_attrs.add((attr.edge_type, attr.layout))
             out.extend(edge_attrs)
+
+        # Check edges added through regular interface:
+        # TODO deprecate this and store edge attributes for all edges in
+        # EdgeStorage
+        for edge_type, edge_store in self.edge_items():
+            for layout, attr_name in EDGE_LAYOUT_TO_ATTR_NAME.items():
+                # Don't double count:
+                if attr_name in edge_store and ((edge_type, layout)
+                                                not in added_attrs):
+                    out.append(
+                        EdgeAttr(edge_type=edge_type, layout=layout,
+                                 size=self[edge_type].size()))
+
         return out
 
 

torch_geometric/loader/neighbor_loader.py

@@ -107,8 +107,8 @@ def __init__(
                     attr.index = None
                 time_tensors = feature_store.multi_get_tensor(time_attrs)
                 self.node_time_dict = {
-                    time_attrs[i].group_name: time_tensors[i]
-                    for i in range(len(time_tensors))
+                    time_attr.group_name: time_tensor
+                    for time_attr, time_tensor in zip(time_attrs, time_tensors)
                 }
 
             # Obtain all node and edge metadata: