precommit

Author: Zack Aristei

test/sampler/test_sampler_base.py

@@ -6,7 +6,7 @@
     NumNeighbors,
     SamplerOutput,
 )
-from torch_geometric.sampler.utils import local_to_global_node_idx, global_to_local_node_idx
+from torch_geometric.sampler.utils import global_to_local_node_idx
 from torch_geometric.testing import get_random_edge_index
 from torch_geometric.utils import is_undirected
 
@@ -121,6 +121,7 @@ def test_heterogeneous_to_bidirectional():
     assert is_undirected(
         torch.stack([obj.row['v1', 'to', 'v1'], obj.col['v1', 'to', 'v1']], 0))
 
+
 def test_homogeneous_sampler_output_global_fields():
     output = SamplerOutput(
         node=torch.tensor([0, 2, 3]),
@@ -152,7 +153,8 @@ def test_homogeneous_sampler_output_global_fields():
     global_values.append(seed_node)
 
     output_bidirectional = output.to_bidirectional(keep_orig_edges=True)
-    global_bidir_row, global_bidir_col = output_bidirectional.global_row, output_bidirectional.global_col
+    global_bidir_row, global_bidir_col = \
+        output_bidirectional.global_row, output_bidirectional.global_col
     assert torch.equal(global_bidir_row, torch.tensor([2, 0, 3, 2]))
     assert torch.equal(global_bidir_col, torch.tensor([0, 2, 2, 3]))
     local_values.append(output_bidirectional.row)
@@ -162,10 +164,12 @@ def test_homogeneous_sampler_output_global_fields():
 
     assert torch.equal(output.global_row, output_bidirectional.global_orig_row)
     assert torch.equal(output.global_col, output_bidirectional.global_orig_col)
-    
+
     # Make sure reverse mapping is correct
     for local_value, global_value in zip(local_values, global_values):
-        assert torch.equal(global_to_local_node_idx(output.node, global_value), local_value)
+        assert torch.equal(global_to_local_node_idx(output.node, global_value),
+                           local_value)
+
 
 def test_heterogeneous_sampler_output_global_fields():
     def _tensor_dict_equal(dict1, dict2):
@@ -177,46 +181,89 @@ def _tensor_dict_equal(dict1, dict2):
 
     output = HeteroSamplerOutput(
         node={"person": torch.tensor([0, 2, 3])},
-        row={("person", "works_with", "person"): torch.tensor([1]), ("person", "leads", "person"): torch.tensor([0])},
-        col={("person", "works_with", "person"): torch.tensor([2]), ("person", "leads", "person"): torch.tensor([1])},
-        edge={("person", "works_with", "person"): torch.tensor([1]), ("person", "leads", "person"): torch.tensor([0])},
+        row={
+            ("person", "works_with", "person"): torch.tensor([1]),
+            ("person", "leads", "person"): torch.tensor([0])
+        },
+        col={
+            ("person", "works_with", "person"): torch.tensor([2]),
+            ("person", "leads", "person"): torch.tensor([1])
+        },
+        edge={
+            ("person", "works_with", "person"): torch.tensor([1]),
+            ("person", "leads", "person"): torch.tensor([0])
+        },
         batch={"person": torch.tensor([0, 0, 0])},
         num_sampled_nodes={"person": torch.tensor([1, 1, 1])},
-        num_sampled_edges={("person", "works_with", "person"): torch.tensor([1]), ("person", "leads", "person"): torch.tensor([1])},
+        num_sampled_edges={
+            ("person", "works_with", "person"): torch.tensor([1]),
+            ("person", "leads", "person"): torch.tensor([1])
+        },
         orig_row=None,
         orig_col=None,
         metadata=(None, None),
     )
 
-    local_values = []
-    global_values = []
-
     global_row, global_col = output.global_row, output.global_col
-    assert _tensor_dict_equal(global_row, {("person", "works_with", "person"): torch.tensor([2]), ("person", "leads", "person"): torch.tensor([0])})
-    assert _tensor_dict_equal(global_col, {("person", "works_with", "person"): torch.tensor([3]), ("person", "leads", "person"): torch.tensor([2])})
-
-    local_row_dict = {k: global_to_local_node_idx(output.node[k[0]], v) for k, v in global_row.items()}
+    assert _tensor_dict_equal(
+        global_row, {
+            ("person", "works_with", "person"): torch.tensor([2]),
+            ("person", "leads", "person"): torch.tensor([0])
+        })
+    assert _tensor_dict_equal(
+        global_col, {
+            ("person", "works_with", "person"): torch.tensor([3]),
+            ("person", "leads", "person"): torch.tensor([2])
+        })
+
+    local_row_dict = {
+        k: global_to_local_node_idx(output.node[k[0]], v)
+        for k, v in global_row.items()
+    }
     assert _tensor_dict_equal(local_row_dict, output.row)
 
-    local_col_dict = {k: global_to_local_node_idx(output.node[k[2]], v) for k, v in global_col.items()}
+    local_col_dict = {
+        k: global_to_local_node_idx(output.node[k[2]], v)
+        for k, v in global_col.items()
+    }
     assert _tensor_dict_equal(local_col_dict, output.col)
 
     seed_node = output.seed_node
     assert _tensor_dict_equal(seed_node, {"person": torch.tensor([0, 0, 0])})
 
-    local_batch_dict = {k: global_to_local_node_idx(output.node[k], v) for k, v in seed_node.items()}
+    local_batch_dict = {
+        k: global_to_local_node_idx(output.node[k], v)
+        for k, v in seed_node.items()
+    }
     assert _tensor_dict_equal(local_batch_dict, output.batch)
 
     output_bidirectional = output.to_bidirectional(keep_orig_edges=True)
-    global_bidir_row, global_bidir_col = output_bidirectional.global_row, output_bidirectional.global_col
-    assert _tensor_dict_equal(global_bidir_row, {("person", "works_with", "person"): torch.tensor([3, 2]), ("person", "leads", "person"): torch.tensor([2, 0])})
-    assert _tensor_dict_equal(global_bidir_col, {("person", "works_with", "person"): torch.tensor([2, 3]), ("person", "leads", "person"): torch.tensor([0, 2])})
-
-    local_bidir_row_dict = {k: global_to_local_node_idx(output_bidirectional.node[k[0]], v) for k, v in global_bidir_row.items()}
+    global_bidir_row, global_bidir_col = \
+        output_bidirectional.global_row, output_bidirectional.global_col
+    assert _tensor_dict_equal(
+        global_bidir_row, {
+            ("person", "works_with", "person"): torch.tensor([3, 2]),
+            ("person", "leads", "person"): torch.tensor([2, 0])
+        })
+    assert _tensor_dict_equal(
+        global_bidir_col, {
+            ("person", "works_with", "person"): torch.tensor([2, 3]),
+            ("person", "leads", "person"): torch.tensor([0, 2])
+        })
+
+    local_bidir_row_dict = {
+        k: global_to_local_node_idx(output_bidirectional.node[k[0]], v)
+        for k, v in global_bidir_row.items()
+    }
     assert _tensor_dict_equal(local_bidir_row_dict, output_bidirectional.row)
 
-    local_bidir_col_dict = {k: global_to_local_node_idx(output_bidirectional.node[k[2]], v) for k, v in global_bidir_col.items()}
+    local_bidir_col_dict = {
+        k: global_to_local_node_idx(output_bidirectional.node[k[2]], v)
+        for k, v in global_bidir_col.items()
+    }
     assert _tensor_dict_equal(local_bidir_col_dict, output_bidirectional.col)
 
-    assert _tensor_dict_equal(output.global_row, output_bidirectional.global_orig_row)
-    assert _tensor_dict_equal(output.global_col, output_bidirectional.global_orig_col)
+    assert _tensor_dict_equal(output.global_row,
+                              output_bidirectional.global_orig_row)
+    assert _tensor_dict_equal(output.global_col,
+                              output_bidirectional.global_orig_col)

torch_geometric/sampler/base.py

@@ -11,7 +11,10 @@
 from torch import Tensor
 
 from torch_geometric.data import Data, FeatureStore, GraphStore, HeteroData
-from torch_geometric.sampler.utils import to_bidirectional, local_to_global_node_idx
+from torch_geometric.sampler.utils import (
+    local_to_global_node_idx,
+    to_bidirectional,
+)
 from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType, OptTensor
 from torch_geometric.utils.mixin import CastMixin
 
@@ -214,19 +217,22 @@ def global_row(self) -> Tensor:
     @property
     def global_col(self) -> Tensor:
         return local_to_global_node_idx(self.node, self.col)
-    
+
     @property
     def seed_node(self) -> Tensor:
-        return local_to_global_node_idx(self.node, self.batch) if self.batch is not None else None
-    
+        return local_to_global_node_idx(
+            self.node, self.batch) if self.batch is not None else None
+
     @property
     def global_orig_row(self) -> Tensor:
-        return local_to_global_node_idx(self.node, self.orig_row) if self.orig_row is not None else None
-    
+        return local_to_global_node_idx(
+            self.node, self.orig_row) if self.orig_row is not None else None
+
     @property
     def global_orig_col(self) -> Tensor:
-        return local_to_global_node_idx(self.node, self.orig_col) if self.orig_col is not None else None
-    
+        return local_to_global_node_idx(
+            self.node, self.orig_col) if self.orig_col is not None else None
+
     def to_bidirectional(
         self,
         keep_orig_edges: bool = False,
@@ -316,23 +322,40 @@ class HeteroSamplerOutput(CastMixin):
 
     @property
     def global_row(self) -> Tensor:
-        return {edge_type: local_to_global_node_idx(self.node[edge_type[0]], row) for edge_type, row in self.row.items()}
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[0]], row)
+            for edge_type, row in self.row.items()
+        }
 
     @property
     def global_col(self) -> Tensor:
-        return {edge_type: local_to_global_node_idx(self.node[edge_type[2]], col) for edge_type, col in self.col.items()}
-    
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[2]], col)
+            for edge_type, col in self.col.items()
+        }
+
     @property
     def seed_node(self) -> Tensor:
-        return {node_type: local_to_global_node_idx(self.node[node_type], batch) for node_type, batch in self.batch.items()}
-    
+        return {
+            node_type: local_to_global_node_idx(self.node[node_type], batch)
+            for node_type, batch in self.batch.items()
+        }
+
     @property
     def global_orig_row(self) -> Tensor:
-        return {edge_type: local_to_global_node_idx(self.node[edge_type[0]], orig_row) for edge_type, orig_row in self.orig_row.items()}
-    
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[0]],
+                                                orig_row)
+            for edge_type, orig_row in self.orig_row.items()
+        }
+
     @property
     def global_orig_col(self) -> Tensor:
-        return {edge_type: local_to_global_node_idx(self.node[edge_type[2]], orig_col) for edge_type, orig_col in self.orig_col.items()}
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[2]],
+                                                orig_col)
+            for edge_type, orig_col in self.orig_col.items()
+        }
 
     def to_bidirectional(
         self,

torch_geometric/sampler/utils.py

@@ -162,35 +162,46 @@ def remap_keys(
     }
 
 
-def local_to_global_node_idx(node_values: Tensor, local_indices: Tensor) -> Tensor:
-    """Convert a tensor of indices referring to elements in the node_values tensor to their values.
+def local_to_global_node_idx(node_values: Tensor,
+                             local_indices: Tensor) -> Tensor:
+    """Convert a tensor of indices referring to elements in the node_values
+    tensor to their values.
 
     Args:
         node_values (Tensor): The node values. (num_nodes, feature_dim)
         local_indices (Tensor): The local indices. (num_indices)
 
     Returns:
-        Tensor: The values of the node_values tensor at the local indices. (num_indices, feature_dim)
+        Tensor: The values of the node_values tensor at the local indices.
+        (num_indices, feature_dim)
     """
     return torch.index_select(node_values, dim=0, index=local_indices)
 
-def global_to_local_node_idx(node_values: Tensor, local_values: Tensor) -> Tensor:
-    """Converts a tensor of values that are contained in the node_values tensor to their indices in that tensor.
+
+def global_to_local_node_idx(node_values: Tensor,
+                             local_values: Tensor) -> Tensor:
+    """Converts a tensor of values that are contained in the node_values
+    tensor to their indices in that tensor.
 
     Args:
         node_values (Tensor): The node values. (num_nodes, feature_dim)
         local_values (Tensor): The local values. (num_indices, feature_dim)
 
     Returns:
-        Tensor: The indices of the local values in the node_values tensor. (num_indices)
+        Tensor: The indices of the local values in the node_values tensor.
+        (num_indices)
     """
     if node_values.dim() == 1:
         node_values = node_values.unsqueeze(1)
     if local_values.dim() == 1:
         local_values = local_values.unsqueeze(1)
-    node_values_expand = node_values.unsqueeze(-1).expand(*node_values.shape, local_values.shape[0]) # (num_nodes, feature_dim, num_indices)
-    local_values_expand = local_values.transpose(0, 1).unsqueeze(0).expand(*node_values_expand.shape) # (num_nodes, feature_dim, num_indices)
-    idx_match = torch.all(node_values_expand == local_values_expand, dim=1).nonzero() # (num_indices, 2)
+    node_values_expand = node_values.unsqueeze(-1).expand(
+        *node_values.shape,
+        local_values.shape[0])  # (num_nodes, feature_dim, num_indices)
+    local_values_expand = local_values.transpose(0, 1).unsqueeze(0).expand(
+        *node_values_expand.shape)  # (num_nodes, feature_dim, num_indices)
+    idx_match = torch.all(node_values_expand == local_values_expand,
+                          dim=1).nonzero()  # (num_indices, 2)
     sort_idx = torch.argsort(idx_match[:, 1])
 
-    return idx_match[:, 0][sort_idx]
+    return idx_match[:, 0][sort_idx]