Repository: pyg-team/pytorch_geometric
Pull request #10315 (2025-06-10 00:34:09)

Sync latest txt2kg with latest bidirectional sampler branch
Changes here will be subsumed when bidirectional sampler branch gets merged into main.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 77cd4a273..e47537c85 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Adds TXT2KG class with examples ([#9992](https://github.com/pyg-team/pytorch_geometric/pull/9992))
+- Added `BidirectionalSampler`, which samples both forwards and backwards on graph edges ([#10126](https://github.com/pyg-team/pytorch_geometric/pull/10126))
+- Enable Sampling both forwards and reverse edges on `NeighborSampler` ([#10126](https://github.com/pyg-team/pytorch_geometric/pull/10126))
+- Added ability to merge together `SamplerOutput` objects ([#10126](https://github.com/pyg-team/pytorch_geometric/pull/10126))
+- Added ability to get global row and col ids from `SamplerOutput` ([#10200](https://github.com/pyg-team/pytorch_geometric/pull/10200))
 - Added PyTorch 2.6 support ([#10170](https://github.com/pyg-team/pytorch_geometric/pull/10170))
 - Added support for heterogenous graphs in `AttentionExplainer` ([#10169](https://github.com/pyg-team/pytorch_geometric/pull/10169))
 - Added support for heterogenous graphs in `PGExplainer` ([#10168](https://github.com/pyg-team/pytorch_geometric/pull/10168))
diff --git a/torch_geometric/sampler/base.py b/torch_geometric/sampler/base.py
index f01e6dfaf..1be1c1e14 100644
--- a/torch_geometric/sampler/base.py
+++ b/torch_geometric/sampler/base.py
@@ -3,7 +3,7 @@ import math
 import warnings
 from abc import ABC
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Union
 
@@ -11,7 +11,12 @@ import torch
 from torch import Tensor
 
 from torch_geometric.data import Data, FeatureStore, GraphStore, HeteroData
-from torch_geometric.sampler.utils import to_bidirectional
+from torch_geometric.sampler.utils import (
+    global_to_local_node_idx,
+    local_to_global_node_idx,
+    to_bidirectional,
+    unique_unsorted,
+)
 from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType, OptTensor
 from torch_geometric.utils.mixin import CastMixin
 
@@ -206,6 +211,39 @@ class SamplerOutput(CastMixin):
     # TODO(manan): refine this further; it does not currently define a proper
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
+    _seed_node: OptTensor = field(repr=False, default=None)
+
+    @property
+    def global_row(self) -> Tensor:
+        return local_to_global_node_idx(self.node, self.row)
+
+    @property
+    def global_col(self) -> Tensor:
+        return local_to_global_node_idx(self.node, self.col)
+
+    @property
+    def seed_node(self) -> Tensor:
+        # can be set manually if the seed nodes are not contained in the
+        # sampled nodes
+        if self._seed_node is None:
+            self._seed_node = local_to_global_node_idx(
+                self.node, self.batch) if self.batch is not None else None
+        return self._seed_node
+
+    @seed_node.setter
+    def seed_node(self, value: Tensor):
+        assert len(value) == len(self.node)
+        self._seed_node = value
+
+    @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
+
+    @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
 
     def to_bidirectional(
         self,
@@ -239,11 +277,13 @@ class SamplerOutput(CastMixin):
 
     @classmethod
     def collate(cls, outputs: List['SamplerOutput'],
-                replace: bool = False) -> 'SamplerOutput':
+                replace: bool = True) -> 'SamplerOutput':
         r"""Collate a list of :class:`~torch_geometric.sampler.SamplerOutput`
         objects into a single :class:`~torch_geometric.sampler.SamplerOutput`
         object. Requires that they all have the same fields.
         """
+        if len(outputs) == 0:
+            raise ValueError("Cannot collate an empty list of SamplerOutputs")
         out = outputs[0]
         has_edge = out.edge is not None
         has_orig_row = out.orig_row is not None
@@ -253,36 +293,37 @@ class SamplerOutput(CastMixin):
         has_num_sampled_edges = out.num_sampled_edges is not None
 
         try:
-            for i, out in enumerate(outputs):
-                assert not has_edge == (out.edge is None)
-                assert not has_orig_row == (out.orig_row is None)
-                assert not has_orig_col == (out.orig_col is None)
-                assert not has_batch == (out.batch is None)
-                assert not has_num_sampled_nodes == (out.num_sampled_nodes
-                                                     is None)
-                assert not has_num_sampled_edges == (out.num_sampled_edges
-                                                     is None)
+            for i, sample_output in enumerate(outputs):
+                assert not has_edge == (sample_output.edge is None)
+                assert not has_orig_row == (sample_output.orig_row is None)
+                assert not has_orig_col == (sample_output.orig_col is None)
+                assert not has_batch == (sample_output.batch is None)
+                assert not has_num_sampled_nodes == (
+                    sample_output.num_sampled_nodes is None)
+                assert not has_num_sampled_edges == (
+                    sample_output.num_sampled_edges is None)
         except AssertionError:
             raise ValueError(
-                f"Output {i} has a different field than the first output")
+                f"Output {i+1} has a different field than the first output")
 
         for other in outputs[1:]:
             out = out.merge_with(other, replace=replace)
         return out
 
     def merge_with(self, other: 'SamplerOutput',
-                   replace: bool = False) -> 'SamplerOutput':
+                   replace: bool = True) -> 'SamplerOutput':
         """Merges two SamplerOutputs.
-        If replace is False, self's nodes and edges take precedence.
+        If replace is True, self's nodes and edges take precedence.
         """
-        if replace:
+        if not replace:
             return SamplerOutput(
                 node=torch.cat([self.node, other.node], dim=0),
                 row=torch.cat([self.row, len(self.node) + other.row], dim=0),
                 col=torch.cat([self.col, len(self.node) + other.col], dim=0),
                 edge=torch.cat([self.edge, other.edge], dim=0)
                 if self.edge is not None and other.edge is not None else None,
-                batch=torch.cat([self.batch, other.batch], dim=0) if
+                batch=torch.cat(
+                    [self.batch, len(self.node) + other.batch], dim=0) if
                 self.batch is not None and other.batch is not None else None,
                 num_sampled_nodes=self.num_sampled_nodes +
                 other.num_sampled_nodes if self.num_sampled_nodes is not None
@@ -310,8 +351,11 @@ class SamplerOutput(CastMixin):
 
             # batch tracks disjoint subgraph samplings
             if self.batch is not None and other.batch is not None:
-                old_node_uid.append(self.batch)
-                new_node_uid.append(other.batch)
+                # Transform the batch indices to be global node ids
+                old_batch_nodes = self.seed_node
+                new_batch_nodes = other.seed_node
+                old_node_uid.append(old_batch_nodes)
+                new_node_uid.append(new_batch_nodes)
 
             # NOTE: if any new node fields are added,
             # they need to be merged here
@@ -319,8 +363,8 @@ class SamplerOutput(CastMixin):
             old_node_uid = torch.stack(old_node_uid, dim=1)
             new_node_uid = torch.stack(new_node_uid, dim=1)
 
-            merged_node_uid = torch.cat([old_node_uid, new_node_uid],
-                                        dim=0).unique(dim=0)
+            merged_node_uid = unique_unsorted(
+                torch.cat([old_node_uid, new_node_uid], dim=0))
             num_old_nodes = old_node_uid.shape[0]
 
             # Recompute num sampled nodes for second output,
@@ -337,13 +381,18 @@ class SamplerOutput(CastMixin):
                     size_of_intersect = torch.cat([
                         old_node_uid,
                         new_node_uid[curr_index:curr_index + minibatch]
-                    ]).unique(dim=0).shape[0] - num_old_nodes
+                    ]).unique(dim=0, sorted=False).shape[0] - num_old_nodes
                     merged_node_num_sampled_nodes.append(size_of_intersect)
                     curr_index += minibatch
 
             merged_nodes = merged_node_uid[:, 0]
-            merged_batch = merged_node_uid[:, 1] \
-                if self.batch is not None and other.batch is not None else None
+            merged_batch = None
+            if self.batch is not None and other.batch is not None:
+                # Restore the batch indices to be relative to the nodes field
+                ref_merged_batch_nodes = merged_node_uid[:, 1].unsqueeze(
+                    -1).expand(-1, 2)  # num_nodes x 2
+                merged_batch = global_to_local_node_idx(
+                    merged_node_uid, ref_merged_batch_nodes)
 
             # EDGES
             is_bidirectional = self.orig_row is not None \
@@ -360,28 +409,54 @@ class SamplerOutput(CastMixin):
             # Transform the row and col indices to be global node ids
             # instead of relative indices to nodes field
             # Edge uids build off of node uids
-            old_row, old_col = torch.index_select(old_node_uid, 0,
-                                                  old_row), torch.index_select(
-                                                      old_node_uid, 0, old_col)
-            new_row, new_col = torch.index_select(new_node_uid, 0,
-                                                  new_row), torch.index_select(
-                                                      new_node_uid, 0, new_col)
+            old_row_idx, old_col_idx = local_to_global_node_idx(
+                old_node_uid,
+                old_row), local_to_global_node_idx(old_node_uid, old_col)
+            new_row_idx, new_col_idx = local_to_global_node_idx(
+                new_node_uid,
+                new_row), local_to_global_node_idx(new_node_uid, new_col)
 
-            old_edge_uid, new_edge_uid = [old_row, old_col], [new_row, new_col]
+            old_edge_uid, new_edge_uid = [old_row_idx, old_col_idx
+                                          ], [new_row_idx, new_col_idx]
 
             row_idx = 0
-            col_idx = old_row.shape[1]
-            edge_idx = old_row.shape[1] + old_col.shape[1]
+            col_idx = old_row_idx.shape[1]
+            edge_idx = old_row_idx.shape[1] + old_col_idx.shape[1]
 
             if self.edge is not None and other.edge is not None:
-                old_edge_uid.append(self.edge.unsqueeze(-1))
-                new_edge_uid.append(other.edge.unsqueeze(-1))
+                if is_bidirectional:
+                    # bidirectional duplicates edge ids
+                    old_edge_uid_ref = torch.stack([self.row, self.col],
+                                                   dim=1)  # num_edges x 2
+                    old_orig_edge_uid_ref = torch.stack(
+                        [self.orig_row, self.orig_col],
+                        dim=1)  # num_orig_edges x 2
+
+                    old_edge_idx = global_to_local_node_idx(
+                        old_edge_uid_ref, old_orig_edge_uid_ref)
+                    old_edge = self.edge[old_edge_idx]
+
+                    new_edge_uid_ref = torch.stack([other.row, other.col],
+                                                   dim=1)  # num_edges x 2
+                    new_orig_edge_uid_ref = torch.stack(
+                        [other.orig_row, other.orig_col],
+                        dim=1)  # num_orig_edges x 2
+
+                    new_edge_idx = global_to_local_node_idx(
+                        new_edge_uid_ref, new_orig_edge_uid_ref)
+                    new_edge = other.edge[new_edge_idx]
+
+                else:
+                    old_edge, new_edge = self.edge, other.edge
+
+                old_edge_uid.append(old_edge.unsqueeze(-1))
+                new_edge_uid.append(new_edge.unsqueeze(-1))
 
             old_edge_uid = torch.cat(old_edge_uid, dim=1)
             new_edge_uid = torch.cat(new_edge_uid, dim=1)
 
-            merged_edge_uid = torch.cat([old_edge_uid, new_edge_uid],
-                                        dim=0).unique(dim=0)
+            merged_edge_uid = unique_unsorted(
+                torch.cat([old_edge_uid, new_edge_uid], dim=0))
             num_old_edges = old_edge_uid.shape[0]
 
             merged_edge_num_sampled_edges = None
@@ -396,30 +471,18 @@ class SamplerOutput(CastMixin):
                     size_of_intersect = torch.cat([
                         old_edge_uid,
                         new_edge_uid[curr_index:curr_index + minibatch]
-                    ]).unique(dim=0).shape[0] - num_old_edges
+                    ]).unique(dim=0, sorted=False).shape[0] - num_old_edges
                     merged_edge_num_sampled_edges.append(size_of_intersect)
                     curr_index += minibatch
 
-            merged_row = merged_edge_uid[:, row_idx]
-            merged_col = merged_edge_uid[:, col_idx]
-            merged_edge = merged_edge_uid[:, edge_idx] \
+            merged_row = merged_edge_uid[:, row_idx:col_idx]
+            merged_col = merged_edge_uid[:, col_idx:edge_idx]
+            merged_edge = merged_edge_uid[:, edge_idx:].squeeze() \
                 if self.edge is not None and other.edge is not None else None
 
             # restore to row and col indices relative to nodes field
-            merged_edge_size = merged_edge.numel()
-
-            merged_node_expand = merged_node_uid.unsqueeze(-1).expand(
-                *merged_node_uid.shape, merged_edge_size)
-
-            merged_row_expand = merged_row.unsqueeze(0).unsqueeze(0).expand(
-                *merged_node_uid.shape, merged_edge_size)
-            merged_row = (merged_row_expand == merged_node_expand).nonzero()[:,
-                                                                             0]
-
-            merged_col_expand = merged_col.unsqueeze(0).unsqueeze(0).expand(
-                *merged_node_uid.shape, merged_edge_size)
-            merged_col = (merged_col_expand == merged_node_expand).nonzero()[:,
-                                                                             0]
+            merged_row = global_to_local_node_idx(merged_node_uid, merged_row)
+            merged_col = global_to_local_node_idx(merged_node_uid, merged_col)
 
             out = SamplerOutput(
                 node=merged_nodes,
@@ -493,6 +556,43 @@ class HeteroSamplerOutput(CastMixin):
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
 
+    @property
+    def global_row(self) -> Dict[EdgeType, Tensor]:
+        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) -> Dict[EdgeType, Tensor]:
+        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) -> Optional[Dict[NodeType, Tensor]]:
+        return {
+            node_type: local_to_global_node_idx(self.node[node_type], batch)
+            for node_type, batch in self.batch.items()
+        } if self.batch is not None else None
+
+    @property
+    def global_orig_row(self) -> Optional[Dict[EdgeType, 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()
+        } if self.orig_row is not None else None
+
+    @property
+    def global_orig_col(self) -> Optional[Dict[EdgeType, 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()
+        } if self.orig_col is not None else None
+
     def to_bidirectional(
         self,
         keep_orig_edges: bool = False,
@@ -575,8 +675,8 @@ class HeteroSamplerOutput(CastMixin):
         return out
 
     @classmethod
-    def collate(cls,
-                outputs: List['HeteroSamplerOutput']) -> 'HeteroSamplerOutput':
+    def collate(cls, outputs: List['HeteroSamplerOutput'],
+                replace: bool = True) -> 'HeteroSamplerOutput':
         r"""Collate a list of
         :class:`~torch_geometric.sampler.HeteroSamplerOutput`objects into a
         single :class:`~torch_geometric.sampler.HeteroSamplerOutput` object.
@@ -586,9 +686,9 @@ class HeteroSamplerOutput(CastMixin):
         raise NotImplementedError
 
     def merge_with(self, other: 'HeteroSamplerOutput',
-                   replace: bool = False) -> 'HeteroSamplerOutput':
+                   replace: bool = True) -> 'HeteroSamplerOutput':
         """Merges two HeteroSamplerOutputs.
-        If replace is False, self's nodes and edges take precedence.
+        If replace is True, self's nodes and edges take precedence.
         """
         # TODO(zaristei)
         raise NotImplementedError
@@ -641,7 +741,7 @@ class NumNeighbors:
             elif isinstance(self.values, dict):
                 default = self.default
             else:
-                raise AssertionError()
+                assert False
 
             # Confirm that `values` only hold valid edge types:
             if isinstance(self.values, dict):
diff --git a/torch_geometric/sampler/neighbor_sampler.py b/torch_geometric/sampler/neighbor_sampler.py
index 2257e6c59..9a8e7b353 100644
--- a/torch_geometric/sampler/neighbor_sampler.py
+++ b/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]]]
@@ -66,6 +72,13 @@ class NeighborSampler(BaseSampler):
         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
 
@@ -101,6 +114,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
@@ -142,7 +166,7 @@ class NeighborSampler(BaseSampler):
                 data, device='cpu', share_memory=share_memory,
                 is_sorted=is_sorted, node_time_dict=self.node_time,
                 edge_time_dict=self.edge_time,
-                to_transpose=sample_direction == 'backward')
+                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)
@@ -173,6 +197,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 "
@@ -269,6 +308,13 @@ class NeighborSampler(BaseSampler):
                     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)
 
@@ -287,18 +333,41 @@ class NeighborSampler(BaseSampler):
         self.subgraph_type = SubgraphType(subgraph_type)
         self.disjoint = disjoint
         self.temporal_strategy = temporal_strategy
-        self.keep_orig_edges = False
 
     @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:
@@ -329,12 +398,8 @@ class NeighborSampler(BaseSampler):
         inputs: NodeSamplerInput,
     ) -> Union[SamplerOutput, HeteroSamplerOutput]:
         out = node_sample(inputs, self._sample)
-        if self.sample_direction == 'backward':
-            row, col = out.row, out.col
-            out.row = col
-            out.col = row
         if self.subgraph_type == SubgraphType.bidirectional:
-            out = out.to_bidirectional(keep_orig_edges=self.keep_orig_edges)
+            out = out.to_bidirectional()
         return out
 
     # Edge-based sampling #####################################################
@@ -346,12 +411,8 @@ class NeighborSampler(BaseSampler):
     ) -> Union[SamplerOutput, HeteroSamplerOutput]:
         out = edge_sample(inputs, self._sample, self.num_nodes, self.disjoint,
                           self.node_time, neg_sampling)
-        if self.sample_direction == 'backward':
-            row, col = out.row, out.col
-            out.row = col
-            out.col = row
         if self.subgraph_type == SubgraphType.bidirectional:
-            out = out.to_bidirectional(keep_orig_edges=self.keep_orig_edges)
+            out = out.to_bidirectional()
         return out
 
     # Other Utilities #########################################################
@@ -448,17 +509,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,
@@ -525,6 +603,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,
@@ -536,15 +617,12 @@ class NeighborSampler(BaseSampler):
             )
 
 
-class BidirectionalNeighborSampler(BaseSampler):
+class BidirectionalNeighborSampler(NeighborSampler):
     """A sampler that allows for both upstream and downstream sampling."""
-
-    # NOTE: Designed to have the same interface as `NeighborSampler`
     def __init__(
         self,
         data: Union[Data, HeteroData, Tuple[FeatureStore, GraphStore]],
         num_neighbors: NumNeighborsType,
-        num_backward_neighbors: Optional[NumNeighborsType] = None,
         subgraph_type: Union[SubgraphType, str] = 'directional',
         replace: bool = False,
         disjoint: bool = False,
@@ -556,38 +634,39 @@ class BidirectionalNeighborSampler(BaseSampler):
         # Deprecated:
         directed: bool = True,
     ):
-        self.backward_neighbors_explicitly_set = (num_backward_neighbors
-                                                  is not None)
-        if num_backward_neighbors is None:
-            num_backward_neighbors = num_neighbors
+
+        # 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_backward_neighbors, subgraph_type, replace, disjoint,
+            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.forward_sampler.num_neighbors
+        return self._num_neighbors
 
     @num_neighbors.setter
     def num_neighbors(self, num_neighbors: NumNeighborsType):
-        self.forward_sampler.num_neighbors = num_neighbors
-        if not self.backward_neighbors_explicitly_set:
-            self.backward_sampler.num_neighbors = num_neighbors
-
-    @property
-    def num_backward_neighbors(self) -> NumNeighbors:
-        return self.backward_sampler.num_neighbors
-
-    @num_backward_neighbors.setter
-    def num_backward_neighbors(self, num_neighbors: NumNeighborsType):
-        self.backward_sampler.num_neighbors = num_neighbors
-        self.backward_neighbors_explicitly_set = True
+        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:
@@ -610,27 +689,105 @@ class BidirectionalNeighborSampler(BaseSampler):
         self,
         inputs: NodeSamplerInput,
     ) -> Union[SamplerOutput, HeteroSamplerOutput]:
-        forward_out = self.forward_sampler.sample_from_nodes(inputs)
-        backward_out = self.backward_sampler.sample_from_nodes(inputs)
-        out = forward_out.merge_with(backward_out)
-        return out
+        return super().sample_from_nodes(inputs)
 
     def sample_from_edges(
         self,
         inputs: EdgeSamplerInput,
         neg_sampling: Optional[NegativeSampling] = None,
     ) -> Union[SamplerOutput, HeteroSamplerOutput]:
-        forward_out = self.forward_sampler.sample_from_edges(
-            inputs, neg_sampling)
-        backward_out = self.backward_sampler.sample_from_edges(
-            inputs, neg_sampling)
-
-        return forward_out.merge_with(backward_out)
+        # 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 ##########################################################
 
diff --git a/torch_geometric/sampler/utils.py b/torch_geometric/sampler/utils.py
index 774dd4a91..184dcb594 100644
--- a/torch_geometric/sampler/utils.py
+++ b/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 ######################################################
 
 
@@ -82,7 +91,8 @@ def to_csc(
 
         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,
@@ -121,6 +131,10 @@ def to_hetero_csc(
         edge_time = (edge_time_dict or {}).get(edge_type, None)
         out = to_csc(store, device, share_memory, is_sorted, src_node_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
@@ -172,3 +186,65 @@ def remap_keys(
         k if k in exclude else mapping.get(k, k): v
         for k, v in inputs.items()
     }
+
+
+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)
+    """
+    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.
+
+    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)
+    """
+    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)
+    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:]))


Test patch:
diff --git a/test/sampler/test_sampler_base.py b/test/sampler/test_sampler_base.py
index 41a7da255..b8d630b90 100644
--- a/test/sampler/test_sampler_base.py
+++ b/test/sampler/test_sampler_base.py
@@ -6,6 +6,7 @@ from torch_geometric.sampler.base import (
     NumNeighbors,
     SamplerOutput,
 )
+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
 
@@ -30,6 +31,168 @@ def test_homogeneous_num_neighbors():
     assert num_neighbors.num_hops == 2  # Test caching.
 
 
+'''
+Merge and collate tests use the following graph:
+
+    #############                    ###########
+    # Alice (0) # -> "works with" -> # Bob (1) #
+    #############                    ###########
+         |
+         v
+      "leads"
+         |
+         v
+    #############                    ############
+    # Carol (2) # -> "works with" -> # Dave (3) #
+    #############                    ############
+
+'''
+
+
+def _init_merge_sampler_outputs(hetero=False, disjoint=False):
+    if not hetero:
+        output1 = SamplerOutput(
+            node=torch.tensor([0, 1, 2]),
+            row=torch.tensor([0, 0]),
+            col=torch.tensor([1, 2]),
+            edge=torch.tensor([0, 1]),
+            batch=torch.tensor([0, 0, 0]) if disjoint else None,
+            num_sampled_nodes=list([1, 2]),
+            num_sampled_edges=list([2]),
+            orig_row=None,
+            orig_col=None,
+            metadata=(None, None),
+        )
+        output2 = SamplerOutput(
+            node=torch.tensor([0, 2, 3]),
+            row=torch.tensor([0, 1]),
+            col=torch.tensor([1, 2]),
+            edge=torch.tensor([1, 2]),
+            batch=torch.tensor([0, 0, 0]) if disjoint else None,
+            num_sampled_nodes=list([1, 1, 1]),
+            num_sampled_edges=list([1, 1]),
+            orig_row=None,
+            orig_col=None,
+            metadata=(None, None),
+        )
+
+        return output1, output2
+    else:
+        # TODO(zaristei)
+        raise NotImplementedError("Heterogeneous merge not implemented")
+
+
+@pytest.mark.parametrize("disjoint", [True, False])
+@pytest.mark.parametrize("bidirectional", [True, False])
+def test_homogeneous_merge(disjoint, bidirectional):
+    """Merge an output representing 1<-0->2 with one representing 0->2->3."""
+    output1, output2 = _init_merge_sampler_outputs(disjoint=disjoint)
+    if bidirectional:
+        output1 = output1.to_bidirectional(keep_orig_edges=True)
+        output2 = output2.to_bidirectional(keep_orig_edges=True)
+
+    expected_output = SamplerOutput(
+        node=torch.tensor([0, 1, 2, 3]),
+        row=torch.tensor([0, 0, 2]),
+        col=torch.tensor([1, 2, 3]),
+        edge=torch.tensor([0, 1, 2]),
+        batch=torch.tensor([0, 0, 0, 0]) if disjoint else None,
+        num_sampled_nodes=[1, 2, 0, 0, 1],
+        num_sampled_edges=[2, 0, 1],
+        orig_row=None,
+        orig_col=None,
+        metadata=[(None, None), (None, None)],
+    )
+    if bidirectional:
+        expected_output = expected_output.to_bidirectional(
+            keep_orig_edges=True)
+    merged_output = output1.merge_with(output2)
+
+    assert str(merged_output) == str(expected_output)
+
+
+@pytest.mark.parametrize("disjoint", [True, False])
+@pytest.mark.parametrize("bidirectional", [True, False])
+def test_homogeneous_merge_no_replace(disjoint, bidirectional):
+    """Merge an output representing 1<-0->2 with one representing 0->2->3.
+    replace=True makes it so that merged output is a simple concatenation
+    instead of removing already sampled nodes/edges.
+    """
+    output1, output2 = _init_merge_sampler_outputs(disjoint=disjoint)
+    if bidirectional:
+        output1 = output1.to_bidirectional(keep_orig_edges=True)
+        output2 = output2.to_bidirectional(keep_orig_edges=True)
+
+    expected_output = SamplerOutput(
+        node=torch.tensor([0, 1, 2, 0, 2, 3]),
+        row=torch.tensor([0, 0, 3, 4]),
+        col=torch.tensor([1, 2, 4, 5]),
+        edge=torch.tensor([0, 1, 1, 2]),
+        batch=torch.tensor([0, 0, 0, 3, 3, 3]) if disjoint else None,
+        num_sampled_nodes=[1, 2, 1, 1, 1],
+        num_sampled_edges=[2, 1, 1],
+        orig_row=None,
+        orig_col=None,
+        metadata=[(None, None), (None, None)],
+    )
+    if bidirectional:
+        expected_output = expected_output.to_bidirectional(
+            keep_orig_edges=True)
+    merged_output = output1.merge_with(output2, replace=False)
+
+    assert str(merged_output) == str(expected_output)
+
+
+def _init_collate_sampler_outputs(disjoint=False):
+    output1, output2 = _init_merge_sampler_outputs(disjoint=disjoint)
+    # new edge not present in graph above
+    output3 = SamplerOutput(
+        node=torch.tensor([3, 4]),
+        row=torch.tensor([0]),
+        col=torch.tensor([1]),
+        edge=torch.tensor([3]),
+        batch=torch.tensor([0, 0]) if disjoint else None,
+        num_sampled_nodes=list([1, 1]),
+        num_sampled_edges=list([1]),
+        orig_row=None,
+        orig_col=None,
+        metadata=(None, None),
+    )
+    return [output1, output2, output3]
+
+
+@pytest.mark.parametrize("replace", [True, False])
+@pytest.mark.parametrize("disjoint", [True, False])
+def test_homogeneous_collate(disjoint, replace):
+    output1, output2, output3 = _init_collate_sampler_outputs(disjoint)
+    collated = SamplerOutput.collate([output1, output2, output3],
+                                     replace=replace)
+    assert str(collated) == str(
+        (output1.merge_with(output2, replace=replace)).merge_with(
+            output3, replace=replace))
+
+
+def test_homogeneous_collate_empty():
+    with pytest.raises(ValueError,
+                       match="Cannot collate an empty list of SamplerOutputs"):
+        SamplerOutput.collate([])
+
+
+def test_homogeneous_collate_single():
+    output, _ = _init_merge_sampler_outputs()
+    collated = SamplerOutput.collate([output])
+    assert str(collated) == str(output)
+
+
+def test_homogeneous_collate_missing_fields():
+    output1, output2, output3 = _init_collate_sampler_outputs()
+    output3.edge = None
+    with pytest.raises(
+            ValueError,
+            match="Output 3 has a different field than the first output"):
+        SamplerOutput.collate([output1, output2, output3])
+
+
 def test_heterogeneous_num_neighbors_list():
     num_neighbors = NumNeighbors([25, 10])
 
@@ -119,3 +282,150 @@ 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]),
+        row=torch.tensor([0, 1]),
+        col=torch.tensor([1, 2]),
+        edge=torch.tensor([1, 2]),
+        batch=torch.tensor([0, 0, 0]),
+        num_sampled_nodes=[1, 1, 1],
+        num_sampled_edges=[1, 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 torch.equal(global_row, torch.tensor([0, 2]))
+    assert torch.equal(global_col, torch.tensor([2, 3]))
+    local_values.append(output.row)
+    local_values.append(output.col)
+    global_values.append(global_row)
+    global_values.append(global_col)
+
+    seed_node = output.seed_node
+    assert torch.equal(seed_node, torch.tensor([0, 0, 0]))
+    local_values.append(output.batch)
+    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
+    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)
+    local_values.append(output_bidirectional.col)
+    global_values.append(global_bidir_row)
+    global_values.append(global_bidir_col)
+
+    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)
+
+
+def test_heterogeneous_sampler_output_global_fields():
+    def _tensor_dict_equal(dict1, dict2):
+        is_equal = True
+        is_equal &= dict1.keys() == dict2.keys()
+        for key in dict1.keys():
+            is_equal &= torch.equal(dict1[key], dict2[key])
+        return is_equal
+
+    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])
+        },
+        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])
+        },
+        orig_row=None,
+        orig_col=None,
+        metadata=(None, None),
+    )
+
+    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(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()
+    }
+    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()
+    }
+    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()
+    }
+    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()
+    }
+    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)
diff --git a/test/sampler/test_sampler_neighbor_sampler.py b/test/sampler/test_sampler_neighbor_sampler.py
index 96cbc63c7..5dedb7e50 100644
--- a/test/sampler/test_sampler_neighbor_sampler.py
+++ b/test/sampler/test_sampler_neighbor_sampler.py
@@ -29,57 +29,106 @@ def _init_sample_graph(hetero=False):
     # Carol (2) # -> "works with" -> # Dave (3) #
     #############                    ############
     """
-    # node attributes dont matter much, they are just necessary for
-    # heterogeneous graphs to work
-    sample_x = torch.tensor([[0], [1], [2], [3]])
-
+    sample_attr = None
+    sample_edge_attr = None
+    sample_edge_indices = None
     if not hetero:
+        sample_attr = torch.tensor([[0], [1], [2], [3]])
+        sample_edge_attr = torch.tensor([[1], [2], [3]])
         sample_edge_indices = torch.tensor([[0, 0, 2], [1, 2, 3]])
     else:
+        sample_attr = dict({
+            "person": dict({"x": torch.tensor([[1], [2], [3]])}),
+            "manager": dict({"x": torch.tensor([[0]])})
+        })
+        sample_edge_attr = dict({
+            ('person', 'works_with', 'person'):
+            dict({"edge_attr": torch.tensor([[3]])}),
+            ('manager', 'leads', 'person'):
+            dict({"edge_attr": torch.tensor([[1]])}),
+            ('manager', 'works_with', 'person'):
+            dict({"edge_attr": torch.tensor([[2]])})
+        })
         sample_edge_indices = dict({
             ('person', 'works_with', 'person'):
-            dict({"edge_index": torch.tensor([[0, 2], [1, 3]])}),
-            ('person', 'leads', 'person'):
-            dict({"edge_index": torch.tensor([[0], [2]])})
+            dict({"edge_index": torch.tensor([[1], [2]])}),
+            ('manager', 'leads', 'person'):
+            dict({"edge_index": torch.tensor([[0], [1]])}),
+            ('manager', 'works_with', 'person'):
+            dict({"edge_index": torch.tensor([[0], [0]])})
         })
-    return sample_x, sample_edge_indices
+    return sample_attr, sample_edge_attr, sample_edge_indices
 
 
 def _init_graph_to_sample(graph_dtype, hetero=False, reverse=False):
-    sample_x, sample_edge_indices = _init_sample_graph(hetero)
+    sample_attr, sample_edge_attr, sample_edge_indices = _init_sample_graph(
+        hetero)
     if reverse:
         if not hetero:
             sample_edge_indices = sample_edge_indices.flip(0)
         else:
+            reversed_edge_indices = dict()
+            reversed_edge_attr = dict()
             for edge_type, edge_index in sample_edge_indices.items():
                 edge_index = edge_index["edge_index"]
+                edge_attr = sample_edge_attr[edge_type]["edge_attr"]
                 flipped_edge_index = edge_index.flip(0)
-                sample_edge_indices[edge_type] = dict(
+                flipped_edge_type = (edge_type[2], edge_type[1], edge_type[0])
+                reversed_edge_indices[flipped_edge_type] = dict(
                     {"edge_index": flipped_edge_index})
+                reversed_edge_attr[flipped_edge_type] = dict(
+                    {"edge_attr": edge_attr})
+            sample_edge_indices = reversed_edge_indices
+            sample_edge_attr = reversed_edge_attr
     graph_to_sample = None
     if graph_dtype == 'data' and not hetero:
-        graph_to_sample = Data(edge_index=sample_edge_indices, x=sample_x)
+        graph_to_sample = Data(edge_index=sample_edge_indices, x=sample_attr,
+                               time=sample_attr.squeeze(-1),
+                               edge_attr=sample_edge_attr.squeeze(-1))
     elif graph_dtype == 'remote' and not hetero:
         graph_store = MyGraphStore()
         graph_store.put_edge_index(sample_edge_indices, edge_type=None,
-                                   layout='coo', size=(4, 4))
+                                   layout='coo', is_sorted=True, size=(4, 4))
         feature_store = MyFeatureStore()
-        feature_store.put_tensor(sample_x, group_name='person', attr_name='x',
+        feature_store.put_tensor(sample_attr, group_name='default',
+                                 attr_name='x', index=None)
+        # temporal node sampling on (fs, gs) needs 'time' attr
+        feature_store.put_tensor(sample_attr.squeeze(-1), group_name='default',
+                                 attr_name='time', index=None)
+        feature_store.put_tensor(sample_edge_attr.squeeze(-1),
+                                 group_name='default', attr_name='edge_attr',
                                  index=None)
         graph_to_sample = (feature_store, graph_store)
     elif graph_dtype == 'data' and hetero:
-        graph_to_sample = HeteroData(sample_edge_indices)
-        graph_to_sample['person'].x = sample_x
+        graph_to_sample = HeteroData()
+        for node_type, node_attr in sample_attr.items():
+            graph_to_sample[node_type].x = node_attr['x']
+            graph_to_sample[node_type].time = node_attr['x'].squeeze(-1)
+        for edge_type in sample_edge_indices.keys():
+            graph_to_sample[edge_type].edge_index = sample_edge_indices[
+                edge_type]["edge_index"]
+            graph_to_sample[edge_type].edge_attr = sample_edge_attr[edge_type][
+                "edge_attr"].squeeze(-1)
     elif graph_dtype == 'remote' and hetero:
         graph_store = MyGraphStore()
         for edge_type, edge_index in sample_edge_indices.items():
             edge_index = edge_index["edge_index"]
-            graph_store.put_edge_index(edge_index, edge_type=edge_type,
-                                       layout='coo',
-                                       size=(len(sample_x), len(sample_x)))
+            graph_store.put_edge_index(
+                edge_index, edge_type=edge_type, layout='coo', is_sorted=True,
+                size=(len(sample_attr[edge_type[0]]["x"]),
+                      len(sample_attr[edge_type[2]]["x"])))
         feature_store = MyFeatureStore()
-        feature_store.put_tensor(sample_x, group_name='person', attr_name='x',
-                                 index=None)
+        for node_type, node_attr in sample_attr.items():
+            feature_store.put_tensor(node_attr["x"], group_name=node_type,
+                                     attr_name='x', index=None)
+            # temporal node sampling on (fs, gs) needs 'time' attr
+            feature_store.put_tensor(node_attr["x"].squeeze(-1),
+                                     group_name=node_type, attr_name='time',
+                                     index=None)
+        for edge_type, edge_attr in sample_edge_attr.items():
+            feature_store.put_tensor(edge_attr["edge_attr"].squeeze(-1),
+                                     group_name=edge_type,
+                                     attr_name='edge_attr', index=None)
         graph_to_sample = (feature_store, graph_store)
     return graph_to_sample
 
@@ -138,33 +187,46 @@ def test_heterogeneous_neighbor_sampler_basic(input_type):
 
     # Sampling from Bob should yield only Alice
     node_sampler_input = NodeSamplerInput(input_id=None,
-                                          node=torch.tensor([1]),
+                                          node=torch.tensor([0]),
                                           input_type="person")
     sampler = NeighborSampler(**sampler_kwargs)
     sampler_output = sampler.sample_from_nodes(node_sampler_input)
-    assert set(sampler_output.node['person'].tolist()) == {1, 0}
-    assert sampler_output.row[('person', 'works_with',
-                               'person')] == torch.tensor([1])
-    assert sampler_output.row[('person', 'leads', 'person')].numel() == 0
-    assert sampler_output.col[('person', 'works_with',
+
+    assert sampler_output.node['person'].tolist() == [0]
+    assert sampler_output.node['manager'].tolist() == [0]
+
+    assert sampler_output.row[('manager', 'works_with',
+                               'person')] == torch.tensor([0])
+    assert sampler_output.row[('manager', 'leads', 'person')].numel() == 0
+    assert sampler_output.row[('person', 'works_with', 'person')].numel() == 0
+    assert sampler_output.col[('manager', 'works_with',
                                'person')] == torch.tensor([0])
-    assert sampler_output.col[('person', 'leads', 'person')].numel() == 0
-    assert sampler_output.edge[('person', 'works_with',
+    assert sampler_output.col[('manager', 'leads', 'person')].numel() == 0
+    assert sampler_output.col[('person', 'works_with', 'person')].numel() == 0
+    assert sampler_output.edge[('manager', 'works_with',
                                 'person')] == torch.tensor([0])
-    assert sampler_output.edge[('person', 'leads', 'person')].numel() == 0
+    assert sampler_output.edge[('manager', 'leads', 'person')].numel() == 0
+    assert sampler_output.edge[('person', 'works_with', 'person')].numel() == 0
 
     # Sampling Alice should yield no edges
     node_sampler_input = NodeSamplerInput(input_id=None,
                                           node=torch.tensor([0]),
-                                          input_type="person")
+                                          input_type="manager")
     sampler_output = sampler.sample_from_nodes(node_sampler_input)
-    assert set(sampler_output.node['person'].tolist()) == {0}
+
+    assert sampler_output.node['manager'].tolist() == [0]
+    assert sampler_output.node['person'].numel() == 0
+
+    assert sampler_output.row[('manager', 'works_with', 'person')].numel() == 0
+    assert sampler_output.row[('manager', 'leads', 'person')].numel() == 0
     assert sampler_output.row[('person', 'works_with', 'person')].numel() == 0
-    assert sampler_output.row[('person', 'leads', 'person')].numel() == 0
+    assert sampler_output.col[('manager', 'works_with', 'person')].numel() == 0
+    assert sampler_output.col[('manager', 'leads', 'person')].numel() == 0
     assert sampler_output.col[('person', 'works_with', 'person')].numel() == 0
-    assert sampler_output.col[('person', 'leads', 'person')].numel() == 0
+    assert sampler_output.edge[('manager', 'works_with',
+                                'person')].numel() == 0
+    assert sampler_output.edge[('manager', 'leads', 'person')].numel() == 0
     assert sampler_output.edge[('person', 'works_with', 'person')].numel() == 0
-    assert sampler_output.edge[('person', 'leads', 'person')].numel() == 0
 
 
 @onlyNeighborSampler
@@ -235,6 +297,115 @@ def test_homogeneous_neighbor_sampler_backwards(input_type):
                        reverse_sampler_output.edge)
 
 
+@onlyNeighborSampler
+@pytest.mark.parametrize('input_type', ['data', 'remote'])
+def test_homogeneous_neighbor_sampler_weighted_backwards(input_type):
+    graph_to_sample = _init_graph_to_sample(input_type, hetero=False)
+    reverse_graph_to_sample = _init_graph_to_sample(input_type, hetero=False,
+                                                    reverse=True)
+
+    sampler_kwargs = {
+        'data': graph_to_sample,
+        'num_neighbors': [1, 1],
+        'weight_attr': 'weight',
+        'sample_direction': 'backward'
+    }
+    reverse_sampler_kwargs = {
+        'data': reverse_graph_to_sample,
+        'num_neighbors': [1, 1],
+        'weight_attr': 'weight',
+        'sample_direction': 'forward',
+    }
+
+    if input_type == 'remote':
+        with pytest.raises(NotImplementedError):
+            NeighborSampler(**sampler_kwargs)
+        return
+
+    graph_to_sample['weight'] = torch.tensor([1.0, 0.0, 1.0])
+    reverse_graph_to_sample['weight'] = torch.tensor([1.0, 0.0, 1.0])
+
+    node_sampler_input = NodeSamplerInput(input_id=None,
+                                          node=torch.tensor([0]))
+
+    # Sampling from Alice should yield Bob
+    backward_sampler = NeighborSampler(**sampler_kwargs)
+    backward_sampler_output = backward_sampler.sample_from_nodes(
+        node_sampler_input)
+
+    reverse_sampler = NeighborSampler(**reverse_sampler_kwargs)
+    reverse_sampler_output = reverse_sampler.sample_from_nodes(
+        node_sampler_input)
+
+    assert torch.equal(backward_sampler_output.node,
+                       reverse_sampler_output.node)
+    assert torch.equal(backward_sampler_output.row, reverse_sampler_output.col)
+    assert torch.equal(backward_sampler_output.col, reverse_sampler_output.row)
+    assert torch.equal(backward_sampler_output.edge,
+                       reverse_sampler_output.edge)
+
+    graph_to_sample['weight'] = torch.tensor([0.0, 1.0, 1.0])
+    reverse_graph_to_sample['weight'] = torch.tensor([0.0, 1.0, 1.0])
+
+    # Sampling from Alice should yield Carol and Dave
+    backward_sampler = NeighborSampler(**sampler_kwargs)
+    backward_sampler_output = backward_sampler.sample_from_nodes(
+        node_sampler_input)
+
+    reverse_sampler = NeighborSampler(**reverse_sampler_kwargs)
+    reverse_sampler_output = reverse_sampler.sample_from_nodes(
+        node_sampler_input)
+
+    assert torch.equal(backward_sampler_output.node,
+                       reverse_sampler_output.node)
+    assert torch.equal(backward_sampler_output.row, reverse_sampler_output.col)
+    assert torch.equal(backward_sampler_output.col, reverse_sampler_output.row)
+    assert torch.equal(backward_sampler_output.edge,
+                       reverse_sampler_output.edge)
+
+
+@onlyNeighborSampler
+@pytest.mark.parametrize('input_type', ['data', 'remote'])
+@pytest.mark.parametrize('time_attr', ['time', 'edge_attr'])
+def test_homogeneous_neighbor_sampler_temporal_backwards(
+        input_type, time_attr):
+    graph_to_sample = _init_graph_to_sample(input_type, hetero=False)
+    reverse_graph_to_sample = _init_graph_to_sample(input_type, hetero=False,
+                                                    reverse=True)
+
+    sampler_kwargs = {
+        'data': graph_to_sample,
+        'num_neighbors': [2, 2],
+        'time_attr': time_attr,
+    }
+    reverse_sampler_kwargs = {
+        'data': reverse_graph_to_sample,
+        'num_neighbors': [2, 2],
+        'time_attr': time_attr,
+    }
+
+    node_sampler_input = NodeSamplerInput(input_id=None,
+                                          node=torch.tensor([1]),
+                                          time=torch.tensor([1]))
+    reverse_node_sampler_input = NodeSamplerInput(input_id=None,
+                                                  node=torch.tensor([0]),
+                                                  time=torch.tensor([1]))
+
+    # sampling from Dave should yield Carol, Alice
+    sampler = NeighborSampler(**sampler_kwargs)
+    sampler_output = sampler.sample_from_nodes(node_sampler_input)
+
+    reverse_sampler = NeighborSampler(**reverse_sampler_kwargs)
+    reverse_sampler_output = reverse_sampler.sample_from_nodes(
+        reverse_node_sampler_input)
+
+    assert torch.equal(sampler_output.node, torch.tensor([1, 0]))
+    assert torch.equal(reverse_sampler_output.node, torch.tensor([0, 1]))
+
+    # TODO (zaristei) Negative cases for temporal sampling, then verify that the output is correct for backwards sampling
+    pytest.skip("still TODO")
+
+
 @onlyNeighborSampler
 @pytest.mark.parametrize('input_type', ['data', 'remote'])
 def test_heterogeneous_neighbor_sampler_backwards(input_type):
@@ -246,7 +417,7 @@ def test_heterogeneous_neighbor_sampler_backwards(input_type):
     }
 
     node_sampler_input = NodeSamplerInput(input_id=None,
-                                          node=torch.tensor([2]),
+                                          node=torch.tensor([1]),
                                           input_type="person")
 
     sampler = NeighborSampler(**sampler_kwargs)
@@ -288,6 +459,9 @@ def test_heterogeneous_neighbor_sampler_backwards(input_type):
     reverse_backward_sampler_output = \
         reverse_backward_sampler.sample_from_nodes(node_sampler_input)
 
+    def reverse_key(key):
+        return (key[2], key[1], key[0])
+
     assert sampler_output.node.keys(
     ) == reverse_backward_sampler_output.node.keys()
     assert reverse_sampler_output.node.keys(
@@ -299,24 +473,32 @@ def test_heterogeneous_neighbor_sampler_backwards(input_type):
         assert torch.equal(reverse_sampler_output.node[key],
                            backward_sampler_output.node[key])
 
-    assert sampler_output.row.keys(
-    ) == reverse_backward_sampler_output.row.keys()
-    assert reverse_sampler_output.row.keys(
-    ) == backward_sampler_output.row.keys()
+    assert len(sampler_output.row.keys()) == len(
+        reverse_backward_sampler_output.row.keys())
     for key in sampler_output.row.keys():
-        assert torch.equal(sampler_output.row[key],
-                           reverse_backward_sampler_output.col[key])
-        assert torch.equal(sampler_output.col[key],
-                           reverse_backward_sampler_output.row[key])
-        assert torch.equal(sampler_output.edge[key],
-                           reverse_backward_sampler_output.edge[key])
+        assert reverse_key(key) in reverse_backward_sampler_output.row.keys()
+
+        assert torch.equal(
+            sampler_output.row[key],
+            reverse_backward_sampler_output.col[reverse_key(key)])
+        assert torch.equal(
+            sampler_output.col[key],
+            reverse_backward_sampler_output.row[reverse_key(key)])
+        assert torch.equal(
+            sampler_output.edge[key],
+            reverse_backward_sampler_output.edge[reverse_key(key)])
+
+    assert len(reverse_sampler_output.row.keys()) == len(
+        backward_sampler_output.row.keys())
     for key in reverse_sampler_output.row.keys():
+        assert reverse_key(key) in backward_sampler_output.row.keys()
+
         assert torch.equal(reverse_sampler_output.row[key],
-                           backward_sampler_output.col[key])
+                           backward_sampler_output.col[reverse_key(key)])
         assert torch.equal(reverse_sampler_output.col[key],
-                           backward_sampler_output.row[key])
+                           backward_sampler_output.row[reverse_key(key)])
         assert torch.equal(reverse_sampler_output.edge[key],
-                           backward_sampler_output.edge[key])
+                           backward_sampler_output.edge[reverse_key(key)])
 
 
 @onlyNeighborSampler
@@ -336,11 +518,11 @@ def test_bidirectional_neighbor_sampler(input_type):
 
     expected_output = SamplerOutput(
         # Union between forward and backward nodes
-        node=torch.tensor([0, 2, 3]),
+        node=torch.tensor([2, 0, 3]),
         # Reindexed to be relative to new nodes field
-        row=torch.tensor([0, 1]),
+        row=torch.tensor([1, 0]),
         # Reindexed to be relative to new nodes field
-        col=torch.tensor([1, 2]),
+        col=torch.tensor([0, 2]),
         # Union between forward and backward edges
         edge=torch.tensor([1, 2]),
         # Will be part of node uid if disjoint=True
@@ -356,9 +538,88 @@ def test_bidirectional_neighbor_sampler(input_type):
         # keep_orig_edges=True
         orig_col=None,
         # simple concat of forward and backward metadata
-        metadata=[(None, None), (None, None)])
+        metadata=(None, None))
     assert str(sampler_output) == str(expected_output)
 
+    adv_sampler_kwargs = {
+        'data': graph_to_sample,
+        'num_neighbors': [2, 2, 2, 2],
+    }
+
+    adv_sampler_input = NodeSamplerInput(input_id=None,
+                                         node=torch.tensor([1, 3]))
+
+    adv_sampler = BidirectionalNeighborSampler(**adv_sampler_kwargs)
+    adv_sampler_output = adv_sampler.sample_from_nodes(adv_sampler_input)
+
+    adv_expected_output = SamplerOutput(
+        node=torch.tensor([1, 3, 0, 2]),
+        row=torch.tensor([2, 3, 2]),
+        col=torch.tensor([0, 1, 3]),
+        edge=torch.tensor([0, 2, 1]),
+        batch=None,
+        # 8 _sample calls total, each have 2 num_sampled_nodes slots
+        num_sampled_nodes=[2, 2] + [0] * 14,
+        num_sampled_edges=[2, 0, 1, 0, 0, 0, 0, 0],
+        orig_row=None,
+        orig_col=None,
+        metadata=(None, None))
+    assert str(adv_sampler_output) == str(adv_expected_output)
+
+    adv_sampler_kwargs['disjoint'] = True
+
+    adv_sampler_disjoint = BidirectionalNeighborSampler(**adv_sampler_kwargs)
+    adv_sampler_disjoint_output = adv_sampler_disjoint.sample_from_nodes(
+        adv_sampler_input)
+
+    adv_expected_disjoint_output = SamplerOutput(
+        node=torch.tensor([1, 3, 0, 2, 0, 2, 1, 3]),
+        row=torch.tensor([2, 3, 4, 2, 4, 5]),
+        col=torch.tensor([0, 1, 3, 5, 6, 7]),
+        edge=torch.tensor([0, 2, 1, 1, 0, 2]),
+        batch=torch.tensor([0, 1, 0, 1, 1, 0, 1, 0]),
+        num_sampled_nodes=[
+            # First forward iteration:
+            # (Bob, Dave) -> (Alice, Carol)
+            2,
+            2,
+            # First backward iteration:
+            # (Bob (seen), Dave (seen)) -> (None, None)
+            0,
+            0,
+            # Second forward iteration:
+            # (Alice (seen), Carol (seen)) -> (None, Alice)
+            0,
+            1,
+            # Second backward iteration:
+            # (Alice (seen), Carol (seen)) ->
+            #   (Bob (seen) and Carol, Alice and Dave (seen))
+            0,
+            1,
+            # Third forward iteration:
+            # (Carol (seen), Alice (seen)) -> (Alice (seen), None)
+            0,
+            0,
+            # Third backward iteration:
+            # (Alice (seen), Carol (seen)) -> (Bob and Carol (seen), Dave)
+            0,
+            2,
+            # Fourth forward iteration:
+            # (Bob (seen), Dave (seen)) -> (Alice (seen), Carol (seen))
+            0,
+            0,
+            # Fourth backward iteration:
+            # (Bob (seen), Dave (seen)) -> (None, None)
+            0,
+            0
+        ],
+        num_sampled_edges=[2, 0, 1, 1, 0, 2, 0, 0],
+        orig_row=None,
+        orig_col=None,
+        metadata=(None, None))
+    assert str(adv_sampler_disjoint_output) == str(
+        adv_expected_disjoint_output)
+
 
 @pytest.mark.skip(
     reason="BidirectionalSampler not implemented yet for heterogeneous graphs."
@@ -367,3 +628,20 @@ def test_bidirectional_neighbor_sampler(input_type):
 @pytest.mark.parametrize('input_type', ['data', 'remote'])
 def test_bidirectional_neighbor_sampler_hetero(input_type):
     raise NotImplementedError
+
+
+@onlyNeighborSampler
+@pytest.mark.parametrize('input_type', ['data', 'remote'])
+@pytest.mark.parametrize('hetero', [False, True])
+def test_neighbor_sampler_backwards_not_supported(input_type, hetero):
+    graph_to_sample = _init_graph_to_sample(input_type, hetero=hetero)
+
+    sampler_kwargs = {
+        'data': graph_to_sample,
+        'num_neighbors': [1],
+        'sample_direction': 'backward',
+        'time_attr': 'time'
+    }
+
+    with pytest.raises(NotImplementedError):
+        NeighborSampler(**sampler_kwargs)