update

Author: rusty1s

test/loader/test_hgt_loader.py

@@ -60,8 +60,9 @@ def test_hgt_loader():
         assert set(batch.node_types) == {'paper', 'author'}
         assert set(batch.edge_types) == set(data.edge_types)
 
-        assert len(batch['paper']) == 2
+        assert len(batch['paper']) == 3
         assert batch['paper'].x.size() == (40, )  # 20 + 4 * 5
+        assert batch['paper'].input_nodes.numel() == batch_size
         assert batch['paper'].batch_size == batch_size
         assert batch['paper'].x.min() >= 0 and batch['paper'].x.max() < 100
 

test/loader/test_link_neighbor_loader.py

@@ -51,9 +51,10 @@ def test_homogeneous_link_neighbor_loader(directed, neg_sampling_ratio):
     for batch in loader:
         assert isinstance(batch, Data)
 
-        assert len(batch) == 5
+        assert len(batch) == 6
         assert batch.x.size(0) <= 100
         assert batch.x.min() >= 0 and batch.x.max() < 100
+        assert batch.input_links.numel() == 20
         assert batch.edge_index.min() >= 0
         assert batch.edge_index.max() < batch.num_nodes
         assert batch.edge_attr.min() >= 0
@@ -110,7 +111,7 @@ def test_heterogeneous_link_neighbor_loader(directed, neg_sampling_ratio):
 
     for batch in loader:
         assert isinstance(batch, HeteroData)
-        assert len(batch) == 5
+        assert len(batch) == 6
         if neg_sampling_ratio == 0.0:
             # Assert only positive samples are present in the original graph:
             assert batch['paper', 'author'].edge_label.sum() == 0
@@ -120,7 +121,6 @@ def test_heterogeneous_link_neighbor_loader(directed, neg_sampling_ratio):
             assert len(edge_index | edge_label_index) == len(edge_index)
 
         else:
-
             assert batch['paper', 'author'].edge_label_index.size(1) == 40
             assert torch.all(batch['paper', 'author'].edge_label[:20] == 1)
             assert torch.all(batch['paper', 'author'].edge_label[20:] == 0)
@@ -312,7 +312,8 @@ def test_homogeneous_link_neighbor_loader_no_edges():
 
     for batch in loader:
         assert isinstance(batch, Data)
-        assert len(batch) == 3
+        assert len(batch) == 4
+        assert batch.input_links.numel() == 20
         assert batch.num_nodes <= 40
         assert batch.edge_label_index.size(1) == 20
         assert batch.num_nodes == batch.edge_label_index.unique().numel()
@@ -328,8 +329,9 @@ def test_heterogeneous_link_neighbor_loader_no_edges():
 
     for batch in loader:
         assert isinstance(batch, HeteroData)
-        assert len(batch) == 3
+        assert len(batch) == 4
         assert batch['paper'].num_nodes <= 40
+        assert batch['paper', 'paper'].input_links.numel() == 20
         assert batch['paper', 'paper'].edge_label_index.size(1) == 20
         assert batch['paper'].num_nodes == batch[
             'paper', 'paper'].edge_label_index.unique().numel()

test/loader/test_neighbor_loader.py

@@ -48,10 +48,9 @@ def test_homogeneous_neighbor_loader(directed):
 
     for batch in loader:
         assert isinstance(batch, Data)
-
-        assert len(batch) == 4
+        assert len(batch) == 5
         assert batch.x.size(0) <= 100
-        assert batch.batch_size == 20
+        assert batch.input_nodes.numel() == batch.batch_size == 20
         assert batch.x.min() >= 0 and batch.x.max() < 100
         assert batch.edge_index.min() >= 0
         assert batch.edge_index.max() < batch.num_nodes
@@ -118,8 +117,9 @@ def test_heterogeneous_neighbor_loader(directed):
         # Test node type selection:
         assert set(batch.node_types) == {'paper', 'author'}
 
-        assert len(batch['paper']) == 2
+        assert len(batch['paper']) == 3
         assert batch['paper'].x.size(0) <= 100
+        assert batch['paper'].input_nodes.numel() == batch_size
         assert batch['paper'].batch_size == batch_size
         assert batch['paper'].x.min() >= 0 and batch['paper'].x.max() < 100
 

torch_geometric/loader/link_loader.py

@@ -135,9 +135,10 @@ def filter_fn(
                                self.link_sampler.edge_permutation)
 
             data.batch = out.batch
-            data.edge_label_index = out.metadata[0]
-            data.edge_label = out.metadata[1]
-            data.edge_label_time = out.metadata[2]
+            data.input_links = out.metadata[0]
+            data.edge_label_index = out.metadata[1]
+            data.edge_label = out.metadata[2]
+            data.edge_label_time = out.metadata[3]
 
         elif isinstance(out, HeteroSamplerOutput):
             if isinstance(self.data, HeteroData):
@@ -150,9 +151,10 @@ def filter_fn(
 
             for key, batch in (out.batch or {}).items():
                 data[key].batch = batch
-            data[self.edge_type].edge_label_index = out.metadata[0]
-            data[self.edge_type].edge_label = out.metadata[1]
-            data[self.edge_type].edge_label_time = out.metadata[2]
+            data[self.edge_type].input_links = out.metadata[0]
+            data[self.edge_type].edge_label_index = out.metadata[1]
+            data[self.edge_type].edge_label = out.metadata[2]
+            data[self.edge_type].edge_label_time = out.metadata[3]
 
         else:
             raise TypeError(f"'{self.__class__.__name__}'' found invalid "

torch_geometric/loader/node_loader.py

@@ -116,7 +116,8 @@ def filter_fn(
             data = filter_data(self.data, out.node, out.row, out.col, out.edge,
                                self.node_sampler.edge_permutation)
             data.batch = out.batch
-            data.batch_size = out.metadata
+            data.input_nodes = out.metadata
+            data.batch_size = out.metadata.size(0)
 
         elif isinstance(out, HeteroSamplerOutput):
             if isinstance(self.data, HeteroData):
@@ -129,7 +130,8 @@ def filter_fn(
 
             for key, batch in (out.batch or {}).items():
                 data[key].batch = batch
-            data[self.node_type].batch_size = out.metadata
+            data[self.node_type].input_nodes = out.metadata
+            data[self.node_type].batch_size = out.metadata.size(0)
 
         else:
             raise TypeError(f"'{self.__class__.__name__}'' found invalid "

torch_geometric/loader/utils.py

@@ -28,7 +28,7 @@ def __getitem__(self, index: Union[Tensor, List[int]]) -> Any:
         if not isinstance(index, Tensor):
             index = torch.tensor(index, dtype=torch.long)
 
-        outs = []
+        outs = [index]
         for arg in self.args:
             outs.append(arg[index] if arg is not None else None)
         return tuple(outs)

torch_geometric/sampler/base.py

@@ -7,16 +7,18 @@
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
 
 # An input to a node-based sampler consists of two tensors:
+#  * The example indices
 #  * The node indices
 #  * The timestamps of the given node indices (optional)
-NodeSamplerInput = Tuple[Tensor, OptTensor]
+NodeSamplerInput = Tuple[Tensor, Tensor, OptTensor]
 
 # An input to an edge-based sampler consists of four tensors:
+#   * The example indices
 #   * The row of the edge index in COO format
 #   * The column of the edge index in COO format
 #   * The labels of the edges
 #   * The time attribute corresponding to the edge label (optional)
-EdgeSamplerInput = Tuple[Tensor, Tensor, Tensor, OptTensor]
+EdgeSamplerInput = Tuple[Tensor, Tensor, Tensor, Tensor, OptTensor]
 
 
 # A sampler output contains the following information.

torch_geometric/sampler/hgt_sampler.py

@@ -60,7 +60,7 @@ def sample_from_nodes(
         index: NodeSamplerInput,
         **kwargs,
     ) -> HeteroSamplerOutput:
-        input_nodes, _ = index
+        index, input_nodes, _ = index
         input_node_dict = {self.input_type: input_nodes}
         sample_fn = torch.ops.torch_sparse.hgt_sample
         out = sample_fn(
@@ -77,7 +77,7 @@ def sample_from_nodes(
             col=remap_keys(col, self.to_edge_type),
             edge=remap_keys(edge, self.to_edge_type),
             batch=batch,
-            metadata=input_nodes.size(0),
+            metadata=index,
         )
 
     def sample_from_edges(

torch_geometric/sampler/neighbor_sampler.py

@@ -326,19 +326,19 @@ def sample_from_nodes(
         index: NodeSamplerInput,
         **kwargs,
     ) -> Union[SamplerOutput, HeteroSamplerOutput]:
-        input_nodes, input_time = index
+        index, input_nodes, input_time = index
 
         if self.data_cls == 'custom' or issubclass(self.data_cls, HeteroData):
             seed_time_dict = None
             if input_time is not None:
                 seed_time_dict = {self.input_type: input_time}
             output = self._sample(seed={self.input_type: input_nodes},
                                   seed_time_dict=seed_time_dict)
-            output.metadata = input_nodes.numel()
+            output.metadata = index
 
         elif issubclass(self.data_cls, Data):
             output = self._sample(seed=input_nodes, seed_time=input_time)
-            output.metadata = input_nodes.numel()
+            output.metadata = index
 
         else:
             raise TypeError(f"'{self.__class__.__name__}'' found invalid "
@@ -353,7 +353,7 @@ def sample_from_edges(
         index: EdgeSamplerInput,
         **kwargs,
     ) -> Union[SamplerOutput, HeteroSamplerOutput]:
-        row, col, edge_label, edge_label_time = index
+        index, row, col, edge_label, edge_label_time = index
         edge_label_index = torch.stack([row, col], dim=0)
         negative_sampling_ratio = kwargs.get('negative_sampling_ratio', 0.0)
 
@@ -421,7 +421,8 @@ def sample_from_edges(
                 for key, batch in output.batch.items():
                     output.batch[key] = batch % num_seed_edges
 
-            output.metadata = (edge_label_index, edge_label, edge_label_time)
+            output.metadata = (index, edge_label_index, edge_label,
+                               edge_label_time)
 
         elif issubclass(self.data_cls, Data):
             if self.disjoint_sampling:
@@ -441,7 +442,8 @@ def sample_from_edges(
             if self.disjoint_sampling:
                 output.batch = output.batch % num_seed_edges
 
-            output.metadata = (edge_label_index, edge_label, edge_label_time)
+            output.metadata = (index, edge_label_index, edge_label,
+                               edge_label_time)
 
         else:
             raise TypeError(f"'{self.__class__.__name__}'' found invalid "