pyg-nightly 2.7.0.dev20250826__py3-none-any.whl → 2.7.0.dev20250827__py3-none-any.whl

torch_geometric/sampler/neighbor_sampler.py

@@ -25,7 +25,13 @@ from torch_geometric.sampler import (
     SamplerOutput,
 )
 from torch_geometric.sampler.base import DataType, NumNeighbors, SubgraphType
-from torch_geometric.sampler.utils import remap_keys, to_csc, to_hetero_csc
+from torch_geometric.sampler.utils import (
+    global_to_local_node_idx,
+    remap_keys,
+    reverse_edge_type,
+    to_csc,
+    to_hetero_csc,
+)
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
 
 NumNeighborsType = Union[NumNeighbors, List[int], Dict[EdgeType, List[int]]]
@@ -47,8 +53,8 @@ class NeighborSampler(BaseSampler):
         weight_attr: Optional[str] = None,
         is_sorted: bool = False,
         share_memory: bool = False,
-        # Deprecated:
-        directed: bool = True,
+        directed: bool = True,  # Deprecated
+        sample_direction: Literal['forward', 'backward'] = 'forward',
     ):
         if not directed:
             subgraph_type = SubgraphType.induced
@@ -66,6 +72,14 @@ class NeighborSampler(BaseSampler):
                 f"accelerated neighborhood sampling", stacklevel=2)
 
         self.data_type = DataType.from_data(data)
+        self.sample_direction = sample_direction
+
+        if self.sample_direction == 'backward':
+            # TODO(zaristei)
+            if time_attr is not None:
+                raise NotImplementedError(
+                    "Temporal Sampling not yet supported for backward sampling"
+                )
 
         if self.data_type == DataType.homogeneous:
             self.num_nodes = data.num_nodes
@@ -87,7 +101,8 @@ class NeighborSampler(BaseSampler):
             self.colptr, self.row, self.perm = to_csc(
                 data, device='cpu', share_memory=share_memory,
                 is_sorted=is_sorted, src_node_time=self.node_time,
-                edge_time=self.edge_time)
+                edge_time=self.edge_time,
+                to_transpose=self.sample_direction == 'backward')
 
             if self.edge_time is not None and self.perm is not None:
                 self.edge_time = self.edge_time[self.perm]
@@ -101,6 +116,17 @@ class NeighborSampler(BaseSampler):
         elif self.data_type == DataType.heterogeneous:
             self.node_types, self.edge_types = data.metadata()
 
+            # reverse edge types if sample_direction is backward
+            if self.sample_direction == 'backward':
+                self.edge_types = [
+                    reverse_edge_type(edge_type)
+                    for edge_type in self.edge_types
+                ]
+                self.to_restored_edge_type = {
+                    k: reverse_edge_type(k)
+                    for k in self.edge_types
+                }
+
             self.num_nodes = {k: data[k].num_nodes for k in self.node_types}
 
             self.node_time: Optional[Dict[NodeType, Tensor]] = None
@@ -141,7 +167,8 @@ class NeighborSampler(BaseSampler):
             colptr_dict, row_dict, self.perm = to_hetero_csc(
                 data, device='cpu', share_memory=share_memory,
                 is_sorted=is_sorted, node_time_dict=self.node_time,
-                edge_time_dict=self.edge_time)
+                edge_time_dict=self.edge_time,
+                to_transpose=self.sample_direction == 'backward')
 
             self.row_dict = remap_keys(row_dict, self.to_rel_type)
             self.colptr_dict = remap_keys(colptr_dict, self.to_rel_type)
@@ -172,6 +199,21 @@ class NeighborSampler(BaseSampler):
             edge_attrs = graph_store.get_all_edge_attrs()
             self.edge_types = list({attr.edge_type for attr in edge_attrs})
 
+            # reverse edge types if sample_direction is backward
+            if self.sample_direction == 'backward':
+                self.edge_types = [
+                    reverse_edge_type(edge_type)
+                    for edge_type in self.edge_types
+                ]
+                self.to_restored_edge_type = {
+                    k: reverse_edge_type(k)
+                    for k in self.edge_types
+                }
+                self.to_backward_edge_type = {
+                    v: k
+                    for k, v in self.to_restored_edge_type.items()
+                }
+
             if weight_attr is not None:
                 raise NotImplementedError(
                     f"'weight_attr' argument not yet supported within "
@@ -221,7 +263,10 @@ class NeighborSampler(BaseSampler):
                     else:
                         self.edge_time = time_tensor
 
-                self.row, self.colptr, self.perm = graph_store.csc()
+                if self.sample_direction == 'forward':
+                    self.row, self.colptr, self.perm = graph_store.csc()
+                elif self.sample_direction == 'backward':
+                    self.colptr, self.row, self.perm = graph_store.csr()
 
             else:
                 node_types = [
@@ -261,8 +306,17 @@ class NeighborSampler(BaseSampler):
                 # Conversion to/from C++ string type (see above):
                 self.to_rel_type = {k: '__'.join(k) for k in self.edge_types}
                 self.to_edge_type = {v: k for k, v in self.to_rel_type.items()}
-                # Convert the graph data into CSC format for sampling:
-                row_dict, colptr_dict, self.perm = graph_store.csc()
+                if self.sample_direction == 'forward':
+                    row_dict, colptr_dict, self.perm = graph_store.csc()
+                elif self.sample_direction == 'backward':
+                    colptr_dict, row_dict, self.perm = graph_store.csr()
+
+                    colptr_dict = remap_keys(colptr_dict,
+                                             self.to_backward_edge_type)
+                    row_dict = remap_keys(row_dict, self.to_backward_edge_type)
+                    self.perm = remap_keys(self.perm,
+                                           self.to_backward_edge_type)
+
                 self.row_dict = remap_keys(row_dict, self.to_rel_type)
                 self.colptr_dict = remap_keys(colptr_dict, self.to_rel_type)
 
@@ -285,14 +339,38 @@ class NeighborSampler(BaseSampler):
 
     @property
     def num_neighbors(self) -> NumNeighbors:
+        if self.sample_direction == 'backward':
+            return self._input_num_neighbors \
+                if self._input_num_neighbors is not None \
+                else self._num_neighbors
         return self._num_neighbors
 
     @num_neighbors.setter
     def num_neighbors(self, num_neighbors: NumNeighborsType):
+        # only used if sample direction is backward and num_neighbors has edge
+        # keys
+        self._input_num_neighbors = None
+
         if isinstance(num_neighbors, NumNeighbors):
-            self._num_neighbors = num_neighbors
+            num_neighbors_values = num_neighbors.values
+            if isinstance(num_neighbors_values,
+                          dict) and self.sample_direction == 'backward':
+                # reverse the edge_types if sample_direction is backward
+                self._input_num_neighbors = num_neighbors
+                num_neighbors_values = remap_keys(num_neighbors_values,
+                                                  self.to_backward_edge_type)
+                self._num_neighbors = NumNeighbors(num_neighbors_values)
+            else:
+                self._num_neighbors = num_neighbors
         else:
-            self._num_neighbors = NumNeighbors(num_neighbors)
+            if isinstance(num_neighbors,
+                          dict) and self.sample_direction == 'backward':
+                # intentionally recursing here to make sure num_neighbors is
+                # set as expected for the user
+                self.num_neighbors = NumNeighbors(
+                    remap_keys(num_neighbors, self.to_backward_edge_type))
+            else:
+                self._num_neighbors = NumNeighbors(num_neighbors)
 
     @property
     def is_hetero(self) -> bool:
@@ -434,17 +512,34 @@ class NeighborSampler(BaseSampler):
                 raise ImportError(f"'{self.__class__.__name__}' requires "
                                   f"either 'pyg-lib' or 'torch-sparse'")
 
+            if self.sample_direction == 'backward':
+                row, col = col, row
+
+            row = remap_keys(row, self.to_edge_type)
+            col = remap_keys(col, self.to_edge_type)
+            edge = remap_keys(edge, self.to_edge_type)
+
+            # In the case of backward sampling, we need to restore the edges
+            # keys to be forward facing in the HeteroSamplerOutput object.
+            if self.sample_direction == 'backward':
+                row = remap_keys(row, self.to_restored_edge_type)
+                col = remap_keys(col, self.to_restored_edge_type)
+                edge = remap_keys(edge, self.to_restored_edge_type)
+
             if num_sampled_edges is not None:
                 num_sampled_edges = remap_keys(
                     num_sampled_edges,
                     self.to_edge_type,
                 )
+                if self.sample_direction == 'backward':
+                    num_sampled_edges = remap_keys(num_sampled_edges,
+                                                   self.to_restored_edge_type)
 
             return HeteroSamplerOutput(
                 node=node,
-                row=remap_keys(row, self.to_edge_type),
-                col=remap_keys(col, self.to_edge_type),
-                edge=remap_keys(edge, self.to_edge_type),
+                row=row,
+                col=col,
+                edge=edge,
                 batch=batch,
                 num_sampled_nodes=num_sampled_nodes,
                 num_sampled_edges=num_sampled_edges,
@@ -511,6 +606,9 @@ class NeighborSampler(BaseSampler):
                 raise ImportError(f"'{self.__class__.__name__}' requires "
                                   f"either 'pyg-lib' or 'torch-sparse'")
 
+            if self.sample_direction == 'backward':
+                row, col = col, row
+
             return SamplerOutput(
                 node=node,
                 row=row,
@@ -522,6 +620,178 @@ class NeighborSampler(BaseSampler):
             )
 
 
+class BidirectionalNeighborSampler(NeighborSampler):
+    """A sampler that allows for both upstream and downstream sampling."""
+    def __init__(
+        self,
+        data: Union[Data, HeteroData, Tuple[FeatureStore, GraphStore]],
+        num_neighbors: NumNeighborsType,
+        subgraph_type: Union[SubgraphType, str] = 'directional',
+        replace: bool = False,
+        disjoint: bool = False,
+        temporal_strategy: str = 'uniform',
+        time_attr: Optional[str] = None,
+        weight_attr: Optional[str] = None,
+        is_sorted: bool = False,
+        share_memory: bool = False,
+        # Deprecated:
+        directed: bool = True,
+    ):
+
+        # TODO(zaristei)
+        if isinstance(num_neighbors, NumNeighbors) and isinstance(
+                num_neighbors.values, dict) or isinstance(num_neighbors, dict):
+            raise RuntimeError(
+                "BidirectionalNeighborSampler does not yet support edge "
+                "delimited sampling.")
+
+        self.forward_sampler = NeighborSampler(
+            data, num_neighbors, subgraph_type, replace, disjoint,
+            temporal_strategy, time_attr, weight_attr, is_sorted, share_memory,
+            sample_direction='forward', directed=directed)
+        self.backward_sampler = NeighborSampler(
+            data, num_neighbors, subgraph_type, replace, disjoint,
+            temporal_strategy, time_attr, weight_attr, is_sorted, share_memory,
+            sample_direction='backward', directed=directed)
+
+        # Trigger warnings on init if number of hops is greater than 1
+        self.num_neighbors = num_neighbors
+        self.subgraph_type = subgraph_type
+
+    @property
+    def num_neighbors(self) -> NumNeighbors:
+        return self._num_neighbors
+
+    @num_neighbors.setter
+    def num_neighbors(self, num_neighbors: NumNeighborsType):
+        if not isinstance(num_neighbors, NumNeighbors):
+            num_neighbors = NumNeighbors(num_neighbors)
+        if num_neighbors.num_hops > 1:
+            print("Warning: Number of hops is greater than 1, resulting in "
+                  "memory-expensive recursive calls.")
+        self._num_neighbors = num_neighbors
+
+    @property
+    def is_hetero(self) -> bool:
+        return self.forward_sampler.is_hetero
+
+    @property
+    def is_temporal(self) -> bool:
+        return self.forward_sampler.is_temporal
+
+    @property
+    def disjoint(self) -> bool:
+        return self.forward_sampler.disjoint
+
+    @disjoint.setter
+    def disjoint(self, disjoint: bool):
+        self.forward_sampler.disjoint = disjoint
+        self.backward_sampler.disjoint = disjoint
+
+    def sample_from_nodes(
+        self,
+        inputs: NodeSamplerInput,
+    ) -> Union[SamplerOutput, HeteroSamplerOutput]:
+        return super().sample_from_nodes(inputs)
+
+    def sample_from_edges(
+        self,
+        inputs: EdgeSamplerInput,
+        neg_sampling: Optional[NegativeSampling] = None,
+    ) -> Union[SamplerOutput, HeteroSamplerOutput]:
+        # TODO(zaristei) Figure out what exactly regular and negative sampling
+        # imply for bidirectional sampling case
+        if neg_sampling is not None:
+            raise RuntimeError(
+                "BidirectionalNeighborSampler does not yet support "
+                "negative sampling.")
+        # Not thoroughly tested yet!
+        return super().sample_from_edges(inputs)
+
+    @property
+    def edge_permutation(self) -> Union[OptTensor, Dict[EdgeType, OptTensor]]:
+        return self.forward_sampler.edge_permutation
+
+    def _sample(
+        self,
+        seed: Union[Tensor, Dict[NodeType, Tensor]],
+        seed_time: Optional[Union[Tensor, Dict[NodeType, Tensor]]] = None,
+        **kwargs,
+    ) -> Union[SamplerOutput, HeteroSamplerOutput]:
+
+        if seed_time is not None:
+            raise NotImplementedError(
+                "BidirectionalNeighborSampler does not yet support "
+                "temporal sampling.")
+
+        if self.is_hetero:
+            raise NotImplementedError(
+                "BidirectionalNeighborSampler does not yet support "
+                "heterogeneous sampling.")
+        else:
+            current_seed = seed
+            current_seed_batch = None
+            current_seed_time = seed_time
+            seen_seed_set = {int(node) for node in current_seed}
+            if self.disjoint:
+                current_seed_batch = torch.arange(len(current_seed))
+                seen_seed_set = {
+                    (int(node), int(batch))
+                    for node, batch in zip(current_seed, current_seed_batch)
+                }
+
+            iter_results = []
+
+            for n_neighbors in self.num_neighbors.values:
+                current_n_neighbors = [n_neighbors]
+                self.forward_sampler.num_neighbors = current_n_neighbors
+                self.backward_sampler.num_neighbors = current_n_neighbors
+
+                fwd_result = self.forward_sampler._sample(
+                    current_seed, current_seed_time, **kwargs)
+                bwd_result = self.backward_sampler._sample(
+                    current_seed, current_seed_time, **kwargs)
+                # The seeds for the next iteration will be the new nodes in
+                # this iteration
+                iter_result = fwd_result.merge_with(bwd_result)
+                iter_results.append(iter_result)
+
+                # Find the nodes not yet seen to set a seed for next iteration
+                if self.disjoint:
+                    iter_seed_global_batch = global_to_local_node_idx(
+                        current_seed_batch, iter_result.batch)
+                    iter_result.seed_node = seed[iter_seed_global_batch]
+
+                    keep_mask = torch.tensor([
+                        (int(node), int(batch)) not in seen_seed_set
+                        for node, batch in zip(iter_result.node,
+                                               iter_seed_global_batch)
+                    ])
+                    next_seed = [(int(node), int(batch))
+                                 for node, batch in zip(
+                                     iter_result.node[keep_mask],
+                                     iter_seed_global_batch[keep_mask])
+                                 ] if keep_mask.any() else []
+                    current_seed, current_seed_batch = torch.tensor(
+                        next_seed).reshape(-1, 2).transpose(0, 1).contiguous()
+                else:
+                    keep_mask = torch.tensor([
+                        int(node) not in seen_seed_set
+                        for node in iter_result.node
+                    ])
+                    next_seed = [
+                        int(node) for node in iter_result.node[keep_mask]
+                    ] if keep_mask.any() else []
+                    current_seed = torch.tensor(next_seed)
+
+                seen_seed_set |= set(next_seed)
+
+                # TODO(zaristei) figure out how to update seed times for
+                # temporal sampling
+
+            return SamplerOutput.collate(iter_results)
+
+
 # Sampling Utilities ##########################################################
 
 

torch_geometric/sampler/utils.py

@@ -9,6 +9,15 @@ from torch_geometric.index import index2ptr
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
 from torch_geometric.utils import coalesce, index_sort, lexsort
 
+
+def reverse_edge_type(edge_type: EdgeType) -> EdgeType:
+    """Reverses edge types for heterogeneous graphs. Useful in cases of
+    backward sampling.
+    """
+    return (edge_type[2], edge_type[1],
+            edge_type[0]) if edge_type is not None else None
+
+
 # Edge Layout Conversion ######################################################
 
 
@@ -41,6 +50,7 @@ def to_csc(
     is_sorted: bool = False,
     src_node_time: Optional[Tensor] = None,
     edge_time: Optional[Tensor] = None,
+    to_transpose: bool = False,
 ) -> Tuple[Tensor, Tensor, OptTensor]:
     # Convert the graph data into a suitable format for sampling (CSC format).
     # Returns the `colptr` and `row` indices of the graph, as well as an
@@ -53,7 +63,10 @@ def to_csc(
         if src_node_time is not None:
             raise NotImplementedError("Temporal sampling via 'SparseTensor' "
                                       "format not yet supported")
-        colptr, row, _ = data.adj.csc()
+        if to_transpose:
+            row, colptr, _ = data.adj.csr()
+        else:
+            colptr, row, _ = data.adj.csc()
 
     elif hasattr(data, 'adj_t'):
         if src_node_time is not None:
@@ -65,13 +78,21 @@ def to_csc(
             # raise NotImplementedError("Temporal sampling via 'SparseTensor' "
             #                           "format not yet supported")
             pass
-        colptr, row, _ = data.adj_t.csr()
+        if to_transpose:
+            row, colptr, _ = data.adj_t.csc()
+        else:
+            colptr, row, _ = data.adj_t.csr()
 
     elif data.edge_index is not None:
-        row, col = data.edge_index
+        if to_transpose:
+            col, row = data.edge_index
+        else:
+            row, col = data.edge_index
+
         if not is_sorted:
             row, col, perm = sort_csc(row, col, src_node_time, edge_time)
-        colptr = index2ptr(col, data.size(1))
+        colptr = index2ptr(col,
+                           data.size(1) if not to_transpose else data.size(0))
     else:
         row = torch.empty(0, dtype=torch.long, device=device)
         colptr = torch.zeros(data.num_nodes + 1, dtype=torch.long,
@@ -97,6 +118,7 @@ def to_hetero_csc(
     is_sorted: bool = False,
     node_time_dict: Optional[Dict[NodeType, Tensor]] = None,
     edge_time_dict: Optional[Dict[EdgeType, Tensor]] = None,
+    to_transpose: bool = False,
 ) -> Tuple[Dict[str, Tensor], Dict[str, Tensor], Dict[str, OptTensor]]:
     # Convert the heterogeneous graph data into a suitable format for sampling
     # (CSC format).
@@ -108,7 +130,11 @@ def to_hetero_csc(
         src_node_time = (node_time_dict or {}).get(edge_type[0], None)
         edge_time = (edge_time_dict or {}).get(edge_type, None)
         out = to_csc(store, device, share_memory, is_sorted, src_node_time,
-                     edge_time)
+                     edge_time, to_transpose)
+        # Edge types need to be reversed for backward sampling:
+        if to_transpose:
+            edge_type = reverse_edge_type(edge_type)
+
         colptr_dict[edge_type], row_dict[edge_type], perm_dict[edge_type] = out
 
     return colptr_dict, row_dict, perm_dict
@@ -205,3 +231,20 @@ def global_to_local_node_idx(node_values: Tensor,
     sort_idx = torch.argsort(idx_match[:, 1])
 
     return idx_match[:, 0][sort_idx]
+
+
+def unique_unsorted(tensor: Tensor) -> Tensor:
+    """Returns the unique elements of a tensor while preserving the original
+    order.
+
+    Necessary because torch.unique() ignores sort parameter.
+    """
+    seen = set()
+    output = []
+    for val in tensor:
+        val = tuple(val.tolist())
+        if val not in seen:
+            seen.add(val)
+            output.append(val)
+    return torch.tensor(output, dtype=tensor.dtype,
+                        device=tensor.device).reshape((-1, *tensor.shape[1:]))