Repository: Snapchat/GiGL
Pull request #22: Add `DistABLPLoader` and transform to convert to homogeneous
Author: login: kmontemayor2-sc
Created: 2025-05-02T20:27:01Z | Merged: 2025-05-14T18:52:52Z

Add `DistABLPLoader` to GiGL as part one of "homogeneous ABLP" support.

Since we want to move towards treating all label edges as full edges in GiGL, we had some request for when there's just one message passing edge type to treat the graph as "homogeneous" and this is the first step towards that.

I will have a follow-up PR to setup the labels properly on the output data object.


Diff:
diff --git a/python/gigl/distributed/distributed_neighborloader.py b/python/gigl/distributed/distributed_neighborloader.py
index e391386b0..74779825a 100644
--- a/python/gigl/distributed/distributed_neighborloader.py
+++ b/python/gigl/distributed/distributed_neighborloader.py
@@ -1,9 +1,11 @@
 from collections import abc
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
+from graphlearn_torch.channel import SampleMessage
 from graphlearn_torch.distributed import DistLoader, MpDistSamplingWorkerOptions
 from graphlearn_torch.sampler import NodeSamplerInput, SamplingConfig, SamplingType
+from torch_geometric.data import Data, HeteroData
 from torch_geometric.typing import EdgeType
 
 import gigl.distributed.utils
@@ -17,6 +19,13 @@
 from gigl.src.common.types.graph_data import (
     NodeType,  # TODO (mkolodner-sc): Change to use torch_geometric.typing
 )
+from gigl.types.graph import (
+    DEFAULT_HOMOGENEOUS_EDGE_TYPE,
+    DEFAULT_HOMOGENEOUS_NODE_TYPE,
+    select_label_edge_types,
+    to_heterogeneous_edge,
+)
+from gigl.utils.data_splitters import get_labels_for_anchor_nodes
 
 logger = Logger()
 
@@ -24,23 +33,67 @@
 DEFAULT_NUM_CPU_THREADS = 2
 
 
-# By default GLT does not support per-example batching, so we add this to support sampling
-# from multiple nodes in a single batch.
-# See https://github.com/alibaba/graphlearn-for-pytorch/pull/155/files for more info.
-class _BatchedNodeSamplerInput(NodeSamplerInput):
+class _LabeledNodeSamplerInput(NodeSamplerInput):
+    """
+    Allows for guaranteeing that all "label sets" are sampled together.
+
+    A "label set" is an anchor node, and it's positive and negative labels.
+    For instance if we have the graph:
+        # Message passing
+        A -> B -> C
+        # Positive label
+        A -> D
+        # Negative label
+        A -> E
+
+    Then the label set for A is (A, D, E).
+
+    If this is being used for labeled input, then the `node` input should be a tensor of shape N x M
+    where N is the number of label sets, and M is the number of nodes in each label set.
+
+    This class may also be given a tensor of shape N x 1, in which case there is no guarantee
+    a node and it's labels are sampled together.
+    """
+
     def __len__(self) -> int:
         return self.node.shape[0]
 
     def __getitem__(
         self, index: Union[torch.Tensor, Any]
-    ) -> "_BatchedNodeSamplerInput":
+    ) -> "_LabeledNodeSamplerInput":
         if not isinstance(index, torch.Tensor):
             index = torch.tensor(index, dtype=torch.long)
         index = index.to(self.node.device)
-        return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
+        return _LabeledNodeSamplerInput(self.node[index].view(-1), self.input_type)
+
+
+class _SupervisedToHomogeneous:
+    """Transform class to convert a heterogeneous graph to a homogeneous graph."""
+
+    def __init__(
+        self,
+        message_passing_edge_type: EdgeType,
+    ):
+        """
+        Args:
+            message_passing_edge_type (EdgeType): The edge type to use for message passing.
+        """
+
+        self._message_passing_edge_type = message_passing_edge_type
+
+    def __call__(self, data: HeteroData) -> Data:
+        """Transform the heterogeneous graph to a homogeneous graph."""
+        homogeneous_data = data.edge_type_subgraph(
+            [self._message_passing_edge_type]
+        ).to_homogeneous(add_edge_type=False, add_node_type=False)
+
+        # TODO(kmonte): Add labels to the returned data.
+        return homogeneous_data
 
 
 class DistNeighborLoader(DistLoader):
+    _transforms: Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]
+
     def __init__(
         self,
         dataset: DistLinkPredictionDataset,
@@ -142,9 +195,20 @@ def __init__(
             # 1. We need to treat `EdgeType` as a proper tuple, not the GiGL`EdgeType`.
             # 2. There are (likely) some GLT bugs around https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_neighbor_sampler.py#L317-L318
             # Where if num_neighbors is a dict then we index into it improperly.
-            raise ValueError(
-                f"num_neighbors must be a list of integers, received: {num_neighbors}"
-            )
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When num_neighbors is a dict, the dataset must be heterogeneous."
+                )
+            if num_neighbors.keys() != dataset.graph.keys():
+                raise ValueError(
+                    f"num_neighbors must have all edge types in the graph, received: {num_neighbors.keys()} with for graph with edge types {dataset.graph.keys()}"
+                )
+            hops = len(next(iter(num_neighbors.values())))
+            if not all(len(fanout) == hops for fanout in num_neighbors.values()):
+                raise ValueError(
+                    f"num_neighbors must be a dict of edge types with the same number of hops. Received: {num_neighbors}"
+                )
+
         curr_process_nodes = _shard_nodes_by_process(
             input_nodes=input_nodes,
             local_process_rank=local_process_rank,
@@ -207,6 +271,19 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # May be set by base classes.
+        if not hasattr(self, "_transforms"):
+            self._transforms = []
+
+        # Determines if the node ids passed in are heterogeneous or homogeneous.
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        input_data = _LabeledNodeSamplerInput(node=node_ids, input_type=node_type)
+
         sampling_config = SamplingConfig(
             sampling_type=SamplingType.NODE,
             num_neighbors=num_neighbors,
@@ -220,15 +297,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
+            )
+        if DEFAULT_HOMOGENEOUS_EDGE_TYPE not in dataset.graph:
+            raise ValueError(
+                f"With Homogeneous ABLP, the graph must have {DEFAULT_HOMOGENEOUS_EDGE_TYPE} edge type, received {dataset.graph.keys()}."
+            )
+        positive_label_edge_type, negative_label_edge_type = select_label_edge_types(
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE, dataset.graph.keys()
+        )
+        positive_labels, negative_labels = get_labels_for_anchor_nodes(
+            dataset, input_nodes, positive_label_edge_type, negative_label_edge_type
+        )
+
+        if negative_labels is not None:
+            input_nodes = torch.cat(
+                [input_nodes.unsqueeze(1), positive_labels, negative_labels], dim=1
+            )
+        else:
+            input_nodes = torch.cat([input_nodes.unsqueeze(1), positive_labels], dim=1)
+        input_nodes = (DEFAULT_HOMOGENEOUS_NODE_TYPE, input_nodes)
+        logger.info(
+            f"Converted input nodes to tuple of ({DEFAULT_HOMOGENEOUS_NODE_TYPE}, {input_nodes[1].shape})."
+        )
+        transforms = [
+            _SupervisedToHomogeneous(
+                message_passing_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
+            )
+        ]
+        self._transforms = transforms
+        # TODO(kmonte): stop setting fanout for positive/negative once GLT sampling is fixed.
+        zero_samples = [0] * len(num_neighbors)
+        num_neighbors = to_heterogeneous_edge(num_neighbors)
+        num_neighbors[positive_label_edge_type] = zero_samples
+        if negative_label_edge_type is not None:
+            num_neighbors[negative_label_edge_type] = zero_samples
+        logger.info(f"Overwrote num_neighbors to: {num_neighbors}.")
+
+        super().__init__(
+            dataset=dataset,
+            num_neighbors=num_neighbors,
+            context=context,
+            local_process_rank=local_process_rank,
+            local_process_world_size=local_process_world_size,
+            input_nodes=input_nodes,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            pin_memory_device=pin_memory_device,
+            worker_concurrency=worker_concurrency,
+            channel_size=channel_size,
+            process_start_gap_seconds=process_start_gap_seconds,
+            num_cpu_threads=num_cpu_threads,
+            shuffle=shuffle,
+            drop_last=drop_last,
+        )
 
 
 def _shard_nodes_by_process(
diff --git a/python/gigl/utils/data_splitters.py b/python/gigl/utils/data_splitters.py
index eb2521513..0a3698f75 100644
--- a/python/gigl/utils/data_splitters.py
+++ b/python/gigl/utils/data_splitters.py
@@ -5,6 +5,7 @@
 
 import torch
 from graphlearn_torch.data import Dataset, Topology
+from torch_geometric.typing import EdgeType as PyGEdgeType
 
 from gigl.common.logger import Logger
 from gigl.src.common.types.graph_data import EdgeType, NodeType
@@ -264,7 +265,8 @@ def __call__(
 def get_labels_for_anchor_nodes(
     dataset: Dataset,
     node_ids: torch.Tensor,
-    supervision_edge_types: Union[EdgeType, tuple[EdgeType, EdgeType]],
+    positive_label_edge_type: PyGEdgeType,
+    negative_label_edge_type: Optional[PyGEdgeType] = None,
 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
     """Selects labels for the given node ids based on the provided edge types.
 
@@ -288,37 +290,25 @@ def get_labels_for_anchor_nodes(
         ]
 
     Args:
-        dataset (Dataset): The dataset storing the graph info.
+        dataset (Dataset): The dataset storing the graph info, must be heterogeneous.
         node_ids (torch.Tensor): The node ids to use for the labels. [N]
-        supervision_edge_types (Union[EdgeType, tuple[EdgeType, EdgeType]]): The edge types to use for the labels.
-            If a single edge type is provided, then the dataset must be heterogeneous,
-            and the edge type provided will be used to generate positive labels.
-            If two edge types are provided, then the dataset must be heterogeneous,
-            the first edge type will be used to generate positive labels, and the second edge type will be used to generate negative labels.
+        positive_label_edge_type (PyGEdgeType): The edge type to use for the positive labels.
+        negative_label_edge_type (Optional[PyGEdgeType]): The edge type to use for the negative labels.
+            Defaults to None. If not provided no negative labels will be returned.
     Returns:
         Tuple of (positive labels, negative_labels?)
-        negative labels may be None depending on supervision_edge_types.
+        negative labels may be None depending on if negative_label_edge_type is provided.
         The returned tensors are of shape N x M where N is the number of nodes and M is the max number of labels, per type.
     """
-    if isinstance(supervision_edge_types, EdgeType):
-        if not isinstance(dataset.graph, Mapping):
-            raise ValueError(
-                f"Got a single supervision edge type {supervision_edge_types} but dataset is homogeneous."
-            )
-        positive_node_topo = dataset.graph[supervision_edge_types].topo
-        negative_node_topo = None
-    elif isinstance(supervision_edge_types, tuple):
-        if not isinstance(dataset.graph, Mapping):
-            raise ValueError(
-                f"Got a single supervision edge type {supervision_edge_types} but dataset is homogeneous."
-            )
-        positive_edge_type, negative_edge_type = supervision_edge_types
-        positive_node_topo = dataset.graph[positive_edge_type].topo
-        negative_node_topo = dataset.graph[negative_edge_type].topo
-    else:
+    if not isinstance(dataset.graph, Mapping):
         raise ValueError(
-            f"Got unknown type for supervision_edge_types: {type(supervision_edge_types)}"
+            "The dataset must be heterogeneous to select labels for anchor nodes."
         )
+    positive_node_topo = dataset.graph[positive_label_edge_type].topo
+    if negative_label_edge_type is not None:
+        negative_node_topo = dataset.graph[negative_label_edge_type].topo
+    else:
+        negative_node_topo = None
 
     # Labels is NxM, where N is the number of nodes, and M is the max number of labels.
     positive_labels = _get_padded_labels(node_ids, positive_node_topo)
diff --git a/python/tests/unit/distributed/distributed_neighborloader_test.py b/python/tests/unit/distributed/distributed_neighborloader_test.py
index 2ff53b2d6..55a04cba9 100644
--- a/python/tests/unit/distributed/distributed_neighborloader_test.py
+++ b/python/tests/unit/distributed/distributed_neighborloader_test.py
@@ -5,12 +5,16 @@
 import graphlearn_torch as glt
 import torch
 import torch.distributed.rpc
+from parameterized import param, parameterized
 from torch.multiprocessing import Manager
 from torch_geometric.data import Data, HeteroData
 
 from gigl.distributed.dist_context import DistributedContext
 from gigl.distributed.dist_link_prediction_dataset import DistLinkPredictionDataset
-from gigl.distributed.distributed_neighborloader import DistNeighborLoader
+from gigl.distributed.distributed_neighborloader import (
+    DistABLPLoader,
+    DistNeighborLoader,
+)
 from gigl.src.common.types.graph_data import NodeType
 from gigl.src.mocking.mocking_assets.mocked_datasets_for_pipeline_tests import (
     CORA_NODE_ANCHOR_MOCKED_DATASET_INFO,
@@ -30,6 +34,9 @@
 )
 from tests.test_assets.distributed.utils import assert_tensor_equality
 
+_POSITIVE_EDGE_TYPE = message_passing_to_positive_label(DEFAULT_HOMOGENEOUS_EDGE_TYPE)
+_NEGATIVE_EDGE_TYPE = message_passing_to_negative_label(DEFAULT_HOMOGENEOUS_EDGE_TYPE)
+
 
 class DistributedNeighborLoaderTest(unittest.TestCase):
     def setUp(self):
@@ -78,12 +85,6 @@ def test_distributed_neighbor_loader(self):
 
     def test_distributed_neighbor_loader_batched(self):
         node_type = DEFAULT_HOMOGENEOUS_NODE_TYPE
-        positive_edge_type = message_passing_to_positive_label(
-            DEFAULT_HOMOGENEOUS_EDGE_TYPE
-        )
-        negative_edge_type = message_passing_to_negative_label(
-            DEFAULT_HOMOGENEOUS_EDGE_TYPE
-        )
         edge_index = {
             DEFAULT_HOMOGENEOUS_EDGE_TYPE: torch.tensor(
                 [
@@ -95,9 +96,8 @@ def test_distributed_neighbor_loader_batched(self):
         partition_output = PartitionOutput(
             node_partition_book=to_heterogeneous_node(torch.zeros(14)),
             edge_partition_book={
-                DEFAULT_HOMOGENEOUS_EDGE_TYPE: torch.zeros(6),
-                positive_edge_type: torch.zeros(3),
-                negative_edge_type: torch.zeros(3),
+                e_type: torch.zeros(e_idx.size(1))
+                for e_type, e_idx in edge_index.items()
             },
             partitioned_edge_index={
                 etype: GraphPartitionData(
@@ -167,6 +167,102 @@ def test_distributed_neighbor_loader_heterogeneous(self):
 
         self.assertEqual(count, 4057)
 
+    @parameterized.expand(
+        [
+            param(
+                "Positive and Negative edges",
+                labeled_edges={
+                    _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]),
+                    _NEGATIVE_EDGE_TYPE: torch.tensor([[10, 11], [16, 14]]),
+                },
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 16, 16, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 12, 14, 13, 17]),
+            ),
+            param(
+                "Positive edges",
+                labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10], [15]])},
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 13, 17]),
+            ),
+        ]
+    )
+    def test_ablp_dataloader(
+        self,
+        _,
+        labeled_edges,
+        expected_node,
+        expected_srcs,
+        expected_dsts,
+    ):
+        # Graph looks like https://is.gd/w2oEVp:
+        # Message passing
+        # 10 -> {11, 12}
+        # 11 -> {13, 17}
+        # 15 -> {13, 14}
+        # 16 -> {12, 14}
+        # Positive labels
+        # 10 -> 15
+        # 15 -> 16
+        # Negative labels
+        # 10 -> 16
+        # 11 -> 14
+
+        edge_index = {
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE: torch.tensor(
+                [
+                    [10, 10, 11, 11, 15, 15, 16, 16],
+                    [11, 12, 13, 17, 13, 14, 12, 14],
+                ]
+            ),
+        }
+        edge_index.update(labeled_edges)
+
+        partition_output = PartitionOutput(
+            node_partition_book=to_heterogeneous_node(torch.zeros(18)),
+            edge_partition_book={
+                e_type: torch.zeros(e_idx.size(1))
+                for e_type, e_idx in edge_index.items()
+            },
+            partitioned_edge_index={
+                etype: GraphPartitionData(
+                    edge_index=idx, edge_ids=torch.arange(idx.size(1))
+                )
+                for etype, idx in edge_index.items()
+            },
+            partitioned_edge_features=None,
+            partitioned_node_features=None,
+            partitioned_negative_labels=None,
+            partitioned_positive_labels=None,
+        )
+        dataset = DistLinkPredictionDataset(rank=0, world_size=1, edge_dir="out")
+        dataset.build(partition_output=partition_output)
+
+        loader = DistABLPLoader(
+            dataset=dataset,
+            num_neighbors=[2, 2],
+            input_nodes=torch.tensor([10]),
+            context=self._context,
+            local_process_rank=0,
+            local_process_world_size=1,
+        )
+
+        count = 0
+        for datum in loader:
+            self.assertIsInstance(datum, Data)
+            count += 1
+
+        self.assertEqual(count, 1)
+        assert_tensor_equality(
+            datum.node,
+            expected_node,
+            dim=0,
+        )
+        dsts, srcs, *_ = datum.coo()
+        assert_tensor_equality(datum.node[srcs], expected_srcs)
+        assert_tensor_equality(datum.node[dsts], expected_dsts)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/python/tests/unit/utils/data_splitters_test.py b/python/tests/unit/utils/data_splitters_test.py
index bd40b4e90..57c938b4c 100644
--- a/python/tests/unit/utils/data_splitters_test.py
+++ b/python/tests/unit/utils/data_splitters_test.py
@@ -536,7 +536,7 @@ def test_get_labels_for_anchor_nodes(self):
         positive, negative = get_labels_for_anchor_nodes(
             dataset=ds,
             node_ids=node_ids,
-            supervision_edge_types=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
+            positive_label_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
         )
         expected_positive = torch.tensor(
             [
@@ -576,7 +576,7 @@ def test_get_labels_for_anchor_nodes_heterogeneous(self):
         # TODO(kmonte): Update to use a splitter once we've migrated splitter API.
         node_ids = torch.tensor([10])
         positive, negative = get_labels_for_anchor_nodes(
-            dataset=ds, node_ids=node_ids, supervision_edge_types=a_to_b
+            dataset=ds, node_ids=node_ids, positive_label_edge_type=a_to_b
         )
         expected = torch.tensor([[10, 11]])
         assert_close(positive, expected, rtol=0, atol=0)
@@ -611,7 +611,8 @@ def test_get_labels_for_anchor_nodes_heterogeneous_positive_and_negative(self):
         positive, negative = get_labels_for_anchor_nodes(
             dataset=ds,
             node_ids=node_ids,
-            supervision_edge_types=(a_to_b, a_to_c),
+            positive_label_edge_type=a_to_b,
+            negative_label_edge_type=a_to_c,
         )
 
         expected_positive = torch.tensor([[10, 11], [12, 13]])


Reviews:
user: login: mkolodner-sc
body: Thanks Kyle, left some comments + questions
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-02T23:28:03Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-05T15:52:33Z
user: login: mkolodner-sc
body: Thanks a lot Kyle! LGTM provided comments are addressed.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-05T18:14:17Z
user: login: yliu2-sc
body: Could you provide more context in the PR description? 
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-05T23:47:21Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T00:26:08Z
user: login: yliu2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T02:15:16Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T22:25:01Z
user: login: mkolodner-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T20:17:19Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T21:37:41Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T21:42:57Z
user: login: mkolodner-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T22:13:04Z
user: login: yliu2-sc
body: Thanks Kyle! I think this structuring makes more sense now.  `to_heterogeneous_edge` and `select_label_edge_types` are not included in this PR btw.
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T22:30:53Z
user: login: mkolodner-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T22:37:06Z
user: login: mkolodner-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T22:38:47Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T23:19:38Z
user: login: mkolodner-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T23:27:11Z
user: login: zfan3-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-10T00:22:01Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T16:59:29Z
user: login: zfan3-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T17:06:51Z
user: login: zfan3-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T17:07:03Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T17:21:43Z
user: login: zfan3-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T17:32:16Z
user: login: zfan3-sc
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T17:32:21Z
user: login: kmontemayor2-sc
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T18:15:26Z
user: login: yliu2-sc
body: The new changes in structure looks fine to me. We should probably rename `_BatchedNodeSamplerInput` --> `_LabeledNodeSamplerInput` or sth more fitting. 
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-13T23:05:55Z

Review comments:
diff_hunk: @@ -200,6 +230,56 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError("Heterogeneous graphs are not supported.")
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: Should this ValueError be suggesting instead that the graph must be heterogeneous in this case, rather than not being supported?
created_at: 2025-05-02T22:54:12Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +230,56 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError("Heterogeneous graphs are not supported.")
+            supervision_edge_types = select_label_edge_types(
+                message_passing_edge_type, dataset.graph.keys()
+            )
+            positive_label_edge_type, negative_label_edge_type = supervision_edge_types
+
+            if supervision_edge_types[1] is None:
+                supervision_edge_types = supervision_edge_types[0]
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: Can we use `positive_label_edge_type ` and `negative_label_edge_type ` edge types here, since they get parsed in the previous line? 

I may have mentioned this before, but it also feels easier to me also to have `get_labels_for_anchor_nodes` accept a `positive_label_edge_type: EdgeType` and `negative_label_edge_type: Optional[EdgeType]` and just be able to pass the parsed items from line 255 directly to `get_labels_for_anchor_nodes` rather than having to modify the length of the passed in Tuple outside of the function. 

For me at least, having an argument of `Union[Tuple[str, str, str], Tuple[Tuple[str, str, str], Tuple[str, str, str]]` feels unecessarily complicated instead of having two arguments: `Tuple[str, str, str]`, `Optional[Tuple[str, str, str]]`. Additionally, inside of `get_labels_for_anchor_nodes`, it's not intuitive to me for why checking if the `len(supervision_edge_types)==3` means we have no negative labels, and `len(supervision_edge_types)==2` means we have negative labels. 
created_at: 2025-05-02T22:56:38Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type (EdgeType): The edge type to use for message passing.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: nit: Optional[EdgeType]

Also, do you think its worth renaming this to `training_message_passing_edge_type`? i see its a bit lengthy, but IIUC we don't want to be setting this if we are doing inference, or else we'll be introducing all the label logic
created_at: 2025-05-02T23:07:48Z
author_association: COLLABORATOR
diff_hunk: @@ -167,6 +164,105 @@ def test_distributed_neighbor_loader_heterogeneous(self):
 
         self.assertEqual(count, 4057)
 
+    @parameterized.expand(
+        [
+            param(
+                "Positive and Negative edges",
+                labeled_edges={
+                    _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]),
+                    _NEGATIVE_EDGE_TYPE: torch.tensor([[10, 11], [16, 14]]),
+                },
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 16, 16, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 12, 14, 13, 17]),
+            ),
+            param(
+                "Positive edges",
+                labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10], [15]])},
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 13, 17]),
+            ),
+        ]
+    )
+    def test_distributed_neighbor_loader_with_supervision_edges(
+        self,
+        _,
+        labeled_edges,
+        expected_node,
+        expected_srcs,
+        expected_dsts,
+    ):
+        node_type = DEFAULT_HOMOGENEOUS_NODE_TYPE
+        # Graph looks like:
+        # Message passing
+        # 10 -> {11, 12}
+        # 11 -> {13, 17}
+        # 15 -> {13, 14}
+        # 16 -> {12, 14}
+        # Positive labels
+        # 10 -> 15
+        # 15 -> 16
+        # Negative labels
+        # 10 -> 16
+        # 11 -> 14
+        # https://dreampuf.github.io/GraphvizOnline/?engine=dot#digraph%20G%20%7B%0A%0A%20%20%20%2010%20-%3E%20%7B11%2C%2012%7D%0A%20%20%20%2011%20-%3E%20%7B13%2C%2017%7D%0A%20%20%20%2015%20-%3E%20%7B13%2C%2014%7D%0A%20%20%20%2016%20-%3E%20%7B12%2C%2014%7D%0A%20%20%20%2010%20-%3E%2015%20%5Bcolor%3D%22blue%22%5D%0A%20%20%20%2015%20-%3E%2016%20%5Bcolor%3D%22blue%22%5D%0A%20%20%20%2010%20-%3E%2016%20%5Bcolor%3D%22red%22%5D%0A%20%20%20%2011%20-%3E%2014%20%5Bcolor%3D%22red%22%5D%0A%7D
path: python/tests/unit/distributed/distributed_neighborloader_test.py
user: login: mkolodner-sc
body: Thanks for this visualization! I think we can use an abbreviated link here -- I followed the link and clicked share -- would https://is.gd/w2oEVp work?
created_at: 2025-05-02T23:20:00Z
author_association: COLLABORATOR
diff_hunk: @@ -167,6 +164,105 @@ def test_distributed_neighbor_loader_heterogeneous(self):
 
         self.assertEqual(count, 4057)
 
+    @parameterized.expand(
+        [
+            param(
+                "Positive and Negative edges",
+                labeled_edges={
+                    _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]),
+                    _NEGATIVE_EDGE_TYPE: torch.tensor([[10, 11], [16, 14]]),
+                },
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 16, 16, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 12, 14, 13, 17]),
path: python/tests/unit/distributed/distributed_neighborloader_test.py
user: login: mkolodner-sc
body: For my learning, why is the order of the src and dsts the same every time?
created_at: 2025-05-02T23:26:43Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +230,56 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError("Heterogeneous graphs are not supported.")
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: oops yeah
created_at: 2025-05-05T15:18:12Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +230,56 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError("Heterogeneous graphs are not supported.")
+            supervision_edge_types = select_label_edge_types(
+                message_passing_edge_type, dataset.graph.keys()
+            )
+            positive_label_edge_type, negative_label_edge_type = supervision_edge_types
+
+            if supervision_edge_types[1] is None:
+                supervision_edge_types = supervision_edge_types[0]
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: sure, updated. 
created_at: 2025-05-05T15:31:28Z
author_association: COLLABORATOR
diff_hunk: @@ -167,6 +164,105 @@ def test_distributed_neighbor_loader_heterogeneous(self):
 
         self.assertEqual(count, 4057)
 
+    @parameterized.expand(
+        [
+            param(
+                "Positive and Negative edges",
+                labeled_edges={
+                    _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]),
+                    _NEGATIVE_EDGE_TYPE: torch.tensor([[10, 11], [16, 14]]),
+                },
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 16, 16, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 12, 14, 13, 17]),
path: python/tests/unit/distributed/distributed_neighborloader_test.py
user: login: kmontemayor2-sc
body: I'm not sure, actually, I'd expect it's because we are always passing in the same nodes to sample from but I'd need to dig more into GLT code to know for sure. 
created_at: 2025-05-05T15:32:28Z
author_association: COLLABORATOR
diff_hunk: @@ -167,6 +164,105 @@ def test_distributed_neighbor_loader_heterogeneous(self):
 
         self.assertEqual(count, 4057)
 
+    @parameterized.expand(
+        [
+            param(
+                "Positive and Negative edges",
+                labeled_edges={
+                    _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]),
+                    _NEGATIVE_EDGE_TYPE: torch.tensor([[10, 11], [16, 14]]),
+                },
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 16, 16, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 12, 14, 13, 17]),
+            ),
+            param(
+                "Positive edges",
+                labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10], [15]])},
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 13, 17]),
+            ),
+        ]
+    )
+    def test_distributed_neighbor_loader_with_supervision_edges(
+        self,
+        _,
+        labeled_edges,
+        expected_node,
+        expected_srcs,
+        expected_dsts,
+    ):
+        node_type = DEFAULT_HOMOGENEOUS_NODE_TYPE
+        # Graph looks like:
+        # Message passing
+        # 10 -> {11, 12}
+        # 11 -> {13, 17}
+        # 15 -> {13, 14}
+        # 16 -> {12, 14}
+        # Positive labels
+        # 10 -> 15
+        # 15 -> 16
+        # Negative labels
+        # 10 -> 16
+        # 11 -> 14
+        # https://dreampuf.github.io/GraphvizOnline/?engine=dot#digraph%20G%20%7B%0A%0A%20%20%20%2010%20-%3E%20%7B11%2C%2012%7D%0A%20%20%20%2011%20-%3E%20%7B13%2C%2017%7D%0A%20%20%20%2015%20-%3E%20%7B13%2C%2014%7D%0A%20%20%20%2016%20-%3E%20%7B12%2C%2014%7D%0A%20%20%20%2010%20-%3E%2015%20%5Bcolor%3D%22blue%22%5D%0A%20%20%20%2015%20-%3E%2016%20%5Bcolor%3D%22blue%22%5D%0A%20%20%20%2010%20-%3E%2016%20%5Bcolor%3D%22red%22%5D%0A%20%20%20%2011%20-%3E%2014%20%5Bcolor%3D%22red%22%5D%0A%7D
path: python/tests/unit/distributed/distributed_neighborloader_test.py
user: login: kmontemayor2-sc
body: seems correct, thanks!
created_at: 2025-05-05T15:35:35Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type (EdgeType): The edge type to use for message passing.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: hmmm, it's not quite `training_message_passing_edge_type` either, right? This this specifically is for auto inferring labels based on some edge type. Maybe `message_passing_edge_type_to_get_labels_from`? Which is even more wordy, but maybe we want to "discourage" people from using this?
created_at: 2025-05-05T15:51:34Z
author_association: COLLABORATOR
diff_hunk: @@ -288,37 +290,25 @@ def get_labels_for_anchor_nodes(
         ]
 
     Args:
-        dataset (Dataset): The dataset storing the graph info.
+        dataset (Dataset): The dataset storing the graph info, must be heterogeneous.
         node_ids (torch.Tensor): The node ids to use for the labels. [N]
-        supervision_edge_types (Union[EdgeType, tuple[EdgeType, EdgeType]]): The edge types to use for the labels.
-            If a single edge type is provided, then the dataset must be heterogeneous,
-            and the edge type provided will be used to generate positive labels.
-            If two edge types are provided, then the dataset must be heterogeneous,
-            the first edge type will be used to generate positive labels, and the second edge type will be used to generate negative labels.
+        positive_label_edge_type (PyGEdgeType): The edge type to use for the positive labels.
+        negative_label_edge_type (Optional[PyGEdgeType]): The edge type to use for the negative labels.
+            Defaults to None. If not provided no negative labels will be returned.
     Returns:
         Tuple of (positive labels, negative_labels?)
-        negative labels may be None depending on supervision_edge_types.
+        negative labels may be None depending on if negative_label_edge_type is provided.
         The returned tensors are of shape N x M where N is the number of nodes and M is the max number of labels, per type.
     """
-    if isinstance(supervision_edge_types, EdgeType):
-        if not isinstance(dataset.graph, Mapping):
-            raise ValueError(
-                f"Got a single supervision edge type {supervision_edge_types} but dataset is homogeneous."
-            )
-        positive_node_topo = dataset.graph[supervision_edge_types].topo
-        negative_node_topo = None
-    elif isinstance(supervision_edge_types, tuple):
-        if not isinstance(dataset.graph, Mapping):
-            raise ValueError(
-                f"Got a single supervision edge type {supervision_edge_types} but dataset is homogeneous."
-            )
-        positive_edge_type, negative_edge_type = supervision_edge_types
-        positive_node_topo = dataset.graph[positive_edge_type].topo
-        negative_node_topo = dataset.graph[negative_edge_type].topo
-    else:
+    if not isinstance(dataset.graph, Mapping):
         raise ValueError(
-            f"Got unknown type for supervision_edge_types: {type(supervision_edge_types)}"
+            "The dataset must be heterogeneous to select labels for anchor nodes."
         )
+    positive_node_topo = dataset.graph[positive_label_edge_type].topo
+    if negative_label_edge_type is not None:
+        negative_node_topo = dataset.graph[negative_label_edge_type].topo
+    else:
+        negative_node_topo = None
path: python/gigl/utils/data_splitters.py
user: login: mkolodner-sc
body: Do we want to do any checks/asserts here to ensure that these edge types are in the graph, and fail appropriately if they are not?
created_at: 2025-05-05T18:07:09Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,10 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type_to_get_labels_from (Optional[EdgeType]): The edge type to use for message passing, only needs to be passed in for training.
+                If passed in, then the dataset must be heterogenneous.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: nit: heterogeneous
created_at: 2025-05-05T18:09:35Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,10 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type_to_get_labels_from (Optional[EdgeType]): The edge type to use for message passing, only needs to be passed in for training.
+                If passed in, then the dataset must be heterogenneous.
+                If passed in, supervision edge types will be automatically selected based on the message passing edge type,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: Should we also add a note that these inferred edge types must be in the `dataset.graph`?
created_at: 2025-05-05T18:13:14Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +44,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _UDLToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: I think we're getting rid of the UDL naming right
created_at: 2025-05-05T22:56:16Z
author_association: COLLABORATOR
diff_hunk: @@ -60,6 +88,7 @@ def __init__(
         num_cpu_threads: Optional[int] = None,
         shuffle: bool = False,
         drop_last: bool = False,
+        message_passing_edge_type_to_get_labels_from: Optional[EdgeType] = None,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: why not just call it `supervision_edge_type`? or `label_edge_type`?
created_at: 2025-05-05T23:25:37Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: add comments that this determines homogeneous vs heterogeneous
created_at: 2025-05-05T23:28:28Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: nit: `message_passing_edge_type` is not a defined variable name
created_at: 2025-05-05T23:29:54Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: Do we only do this for homogeneous cases? 
created_at: 2025-05-05T23:32:54Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
+                message_passing_edge_type_to_get_labels_from, dataset.graph.keys()
+            )
+
+            positive_labels, negative_labels = get_labels_for_anchor_nodes(
+                dataset, node_ids, positive_label_edge_type, negative_label_edge_type
+            )
+            if negative_labels is not None:
+                node_ids = torch.cat(
+                    [node_ids.unsqueeze(1), positive_labels, negative_labels], dim=1
+                )
+            else:
+                node_ids = torch.cat([node_ids.unsqueeze(1), positive_labels], dim=1)
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: can probably just keep these separate and not cat
created_at: 2025-05-05T23:35:11Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
+                message_passing_edge_type_to_get_labels_from, dataset.graph.keys()
+            )
+
+            positive_labels, negative_labels = get_labels_for_anchor_nodes(
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: Is the purpose of the addition here to load labels? In that case this should be part of the ABLP loader and not DistNeighborLoader? 
created_at: 2025-05-05T23:46:40Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +44,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _UDLToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Yeah, WDYT about `_SupervisedToHomogeneous`?
created_at: 2025-05-06T00:03:20Z
author_association: COLLABORATOR
diff_hunk: @@ -60,6 +88,7 @@ def __init__(
         num_cpu_threads: Optional[int] = None,
         shuffle: bool = False,
         drop_last: bool = False,
+        message_passing_edge_type_to_get_labels_from: Optional[EdgeType] = None,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: This is the edge type that the labels are based on, I'm not sure it's "right" to call it a label edge type, this is like `user-to-user` and then we'd have `user-positive-user` and `user-negative-user` which are inferred in the graph. 

I guess this edge type is "supervised" but I feel like if I saw "supervision_edge_type" I'd assume it was the edge type that contains the labels.

WDYT?
created_at: 2025-05-06T00:05:49Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,10 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type_to_get_labels_from (Optional[EdgeType]): The edge type to use for message passing, only needs to be passed in for training.
+                If passed in, then the dataset must be heterogenneous.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: thanks!
created_at: 2025-05-06T00:06:09Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,10 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type_to_get_labels_from (Optional[EdgeType]): The edge type to use for message passing, only needs to be passed in for training.
+                If passed in, then the dataset must be heterogenneous.
+                If passed in, supervision edge types will be automatically selected based on the message passing edge type,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Hmmm, I think that ideally users shouldn't have to worry about this at all, I've left a todo to add more details about how the inferred edge types get there once this is more fleshed out.
created_at: 2025-05-06T00:20:58Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Done.
created_at: 2025-05-06T00:21:41Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: good catch, done 
created_at: 2025-05-06T00:22:36Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: hmmm WDYM homogeneous? We can really only do this if the dataset is heterogeneous, and I guess this should still be valid if there are multiple sets of supervision edge types. 
created_at: 2025-05-06T00:23:35Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
+                message_passing_edge_type_to_get_labels_from, dataset.graph.keys()
+            )
+
+            positive_labels, negative_labels = get_labels_for_anchor_nodes(
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: I wanted to keep us with one loader - but if you feel strongly that we should split up the loaders to be more task specific we can do that.
created_at: 2025-05-06T00:24:01Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
+                message_passing_edge_type_to_get_labels_from, dataset.graph.keys()
+            )
+
+            positive_labels, negative_labels = get_labels_for_anchor_nodes(
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: I don't think that PyG really provides guidance here right, it's just about what we're sampling from, and in this case we are sampling from nodes in either case. 
created_at: 2025-05-06T00:24:38Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
+                message_passing_edge_type_to_get_labels_from, dataset.graph.keys()
+            )
+
+            positive_labels, negative_labels = get_labels_for_anchor_nodes(
+                dataset, node_ids, positive_label_edge_type, negative_label_edge_type
+            )
+            if negative_labels is not None:
+                node_ids = torch.cat(
+                    [node_ids.unsqueeze(1), positive_labels, negative_labels], dim=1
+                )
+            else:
+                node_ids = torch.cat([node_ids.unsqueeze(1), positive_labels], dim=1)
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: In that case we'd do the `cat` in the `_BatchedNodeSamplerInput` instead? We could do that but I'm not sure that's much cleaner?
created_at: 2025-05-06T00:25:36Z
author_association: COLLABORATOR
diff_hunk: @@ -288,37 +290,25 @@ def get_labels_for_anchor_nodes(
         ]
 
     Args:
-        dataset (Dataset): The dataset storing the graph info.
+        dataset (Dataset): The dataset storing the graph info, must be heterogeneous.
         node_ids (torch.Tensor): The node ids to use for the labels. [N]
-        supervision_edge_types (Union[EdgeType, tuple[EdgeType, EdgeType]]): The edge types to use for the labels.
-            If a single edge type is provided, then the dataset must be heterogeneous,
-            and the edge type provided will be used to generate positive labels.
-            If two edge types are provided, then the dataset must be heterogeneous,
-            the first edge type will be used to generate positive labels, and the second edge type will be used to generate negative labels.
+        positive_label_edge_type (PyGEdgeType): The edge type to use for the positive labels.
+        negative_label_edge_type (Optional[PyGEdgeType]): The edge type to use for the negative labels.
+            Defaults to None. If not provided no negative labels will be returned.
     Returns:
         Tuple of (positive labels, negative_labels?)
-        negative labels may be None depending on supervision_edge_types.
+        negative labels may be None depending on if negative_label_edge_type is provided.
         The returned tensors are of shape N x M where N is the number of nodes and M is the max number of labels, per type.
     """
-    if isinstance(supervision_edge_types, EdgeType):
-        if not isinstance(dataset.graph, Mapping):
-            raise ValueError(
-                f"Got a single supervision edge type {supervision_edge_types} but dataset is homogeneous."
-            )
-        positive_node_topo = dataset.graph[supervision_edge_types].topo
-        negative_node_topo = None
-    elif isinstance(supervision_edge_types, tuple):
-        if not isinstance(dataset.graph, Mapping):
-            raise ValueError(
-                f"Got a single supervision edge type {supervision_edge_types} but dataset is homogeneous."
-            )
-        positive_edge_type, negative_edge_type = supervision_edge_types
-        positive_node_topo = dataset.graph[positive_edge_type].topo
-        negative_node_topo = dataset.graph[negative_edge_type].topo
-    else:
+    if not isinstance(dataset.graph, Mapping):
         raise ValueError(
-            f"Got unknown type for supervision_edge_types: {type(supervision_edge_types)}"
+            "The dataset must be heterogeneous to select labels for anchor nodes."
         )
+    positive_node_topo = dataset.graph[positive_label_edge_type].topo
+    if negative_label_edge_type is not None:
+        negative_node_topo = dataset.graph[negative_label_edge_type].topo
+    else:
+        negative_node_topo = None
path: python/gigl/utils/data_splitters.py
user: login: kmontemayor2-sc
body: won't we throw when we do the `[]` access?
created_at: 2025-05-06T00:25:59Z
author_association: COLLABORATOR
diff_hunk: @@ -60,6 +88,7 @@ def __init__(
         num_cpu_threads: Optional[int] = None,
         shuffle: bool = False,
         drop_last: bool = False,
+        message_passing_edge_type_to_get_labels_from: Optional[EdgeType] = None,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: What do we expect users do pass in actually? I'd imagine that for 
- homogeneous cases: users pass in the edge type, say for user-friend-user graph, they'd pass in  user-friend-user as supervision or label edge type, and we deal with it under the hood
- heterogeneous cases: same thing, but for multiple supervision edge types, we'd just have multiple distNeighborLoader

So from the user's perspective no matter homogeneous or heterogeneous, they'd be specifying the edge label or supervision edge type. 
created_at: 2025-05-08T01:56:22Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,11 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type_to_get_labels_from (Optional[EdgeType]): The edge type to use for message passing, only needs to be passed in for training.
+                If passed in, then the dataset must be heterogeneous.
+                If passed in, supervision edge types will be automatically selected based on the message passing edge type,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: I don't think users need to care about this, this homogeneous and heterogeneous is our internal representation and not what the users see, right? 
created_at: 2025-05-08T02:00:06Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
+                message_passing_edge_type_to_get_labels_from, dataset.graph.keys()
+            )
+
+            positive_labels, negative_labels = get_labels_for_anchor_nodes(
+                dataset, node_ids, positive_label_edge_type, negative_label_edge_type
+            )
+            if negative_labels is not None:
+                node_ids = torch.cat(
+                    [node_ids.unsqueeze(1), positive_labels, negative_labels], dim=1
+                )
+            else:
+                node_ids = torch.cat([node_ids.unsqueeze(1), positive_labels], dim=1)
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: Oh I see this is what goes into the sampler, but what about the return from the distNeighborloader? Is it stored all in batch.y? 
created_at: 2025-05-08T02:14:37Z
author_association: COLLABORATOR
diff_hunk: @@ -200,6 +233,62 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
+        self._transforms: list[
+            Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]
+        ] = []
+
+        if isinstance(curr_process_nodes, torch.Tensor):
+            node_ids = curr_process_nodes
+            node_type = None
+        else:
+            node_type, node_ids = curr_process_nodes
+
+        # TODO(kmonte): Once GLT sampling works more generally, we should only sample against message_passing_edge_type for this code path.
+        if message_passing_edge_type_to_get_labels_from is not None:
+            if len(node_ids.shape) != 1:
+                raise ValueError(
+                    f"node_ids must be a 1D tensor when supervision_edge_types are provided, got {node_ids.shape}."
+                )
+
+            if not isinstance(dataset.graph, abc.Mapping):
+                raise ValueError(
+                    "When `message_passing_edge_type` is provided, the dataset must be heterogeneous."
+                )
+            (
+                positive_label_edge_type,
+                negative_label_edge_type,
+            ) = select_label_edge_types(
+                message_passing_edge_type_to_get_labels_from, dataset.graph.keys()
+            )
+
+            positive_labels, negative_labels = get_labels_for_anchor_nodes(
+                dataset, node_ids, positive_label_edge_type, negative_label_edge_type
+            )
+            if negative_labels is not None:
+                node_ids = torch.cat(
+                    [node_ids.unsqueeze(1), positive_labels, negative_labels], dim=1
+                )
+            else:
+                node_ids = torch.cat([node_ids.unsqueeze(1), positive_labels], dim=1)
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: I'm not modifying anything in the returned data in this PR, so `Data.y` will be unset.

All sampled node ids will be in `Data.node` and the "root" node ids will be in `Data.batch`.

See an example output:

```
Data(
  num_sampled_nodes={ default_homogeneous_node_type=[3] },
  num_sampled_edges={
    (default_homogeneous_node_type, to, default_homogeneous_node_type)=[2],
    (default_homogeneous_node_type, to_gigl_positive, default_homogeneous_node_type)=[2],
  },
  edge_index=[2, 6],
  node=[7],
  batch=[2],
  edge=[6]
)
```
created_at: 2025-05-08T17:09:44Z
author_association: COLLABORATOR
diff_hunk: @@ -60,6 +88,7 @@ def __init__(
         num_cpu_threads: Optional[int] = None,
         shuffle: bool = False,
         drop_last: bool = False,
+        message_passing_edge_type_to_get_labels_from: Optional[EdgeType] = None,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Hmmm you know, I think it's better to have this be a separate `DistABLPNeighborLoader`, so we can more strongly differentiate between this case and the more general case.

I think in the ABLP case users should be expected to pass in:

If """"homogenous""" (one message passing edge type, one or two label edge types):

1. The nodes to sample against
2. The dataset.

Then we infer the message passing edge type/node type/label edge types/etc. They never need to be exposed to  our internal edge types/etc.

For truely heterogeneous cases they should be expected to pass in:

1. The `(node_type, node_ids)` tuple
2. The dataset
3. The edge type(s) to sample against.

I can make these changes.
created_at: 2025-05-08T17:25:13Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +138,11 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            message_passing_edge_type_to_get_labels_from (Optional[EdgeType]): The edge type to use for message passing, only needs to be passed in for training.
+                If passed in, then the dataset must be heterogeneous.
+                If passed in, supervision edge types will be automatically selected based on the message passing edge type,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Probably not, hopefully adding `DistNABLPLoader` makes things more clear!
created_at: 2025-05-08T22:23:38Z
author_association: COLLABORATOR
diff_hunk: @@ -207,6 +255,20 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: nit: we can remove this line now
created_at: 2025-05-09T19:59:04Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +145,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            transforms (Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]): A sequence of transforms to apply to the data.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: What transforms do you imagine we'll need to be passing in the future to the `DistNeighborLoader`? Should we make the transform input arg and collate specific to the ABLP dataloader, rather than relying on an inheritance pattern here?
created_at: 2025-05-09T20:01:03Z
author_association: COLLABORATOR
diff_hunk: @@ -167,6 +167,103 @@ def test_distributed_neighbor_loader_heterogeneous(self):
 
         self.assertEqual(count, 4057)
 
+    @parameterized.expand(
+        [
+            param(
+                "Positive and Negative edges",
+                labeled_edges={
+                    _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]),
+                    _NEGATIVE_EDGE_TYPE: torch.tensor([[10, 11], [16, 14]]),
+                },
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 16, 16, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 12, 14, 13, 17]),
+            ),
+            param(
+                "Positive edges",
+                labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10], [15]])},
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 13, 17]),
+            ),
+        ]
+    )
+    def test_distributed_neighbor_loader_with_supervision_edges(
path: python/tests/unit/distributed/distributed_neighborloader_test.py
user: login: mkolodner-sc
body: nit: can we call this `test_ablp_dataloader`?
created_at: 2025-05-09T20:12:14Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +282,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: int = 60,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: Small note: I have a PR https://github.com/Snapchat/GiGL/pull/39/files#diff-e751a9cd787e206721cd05c941effd36bae4e56559ac269889b7e9fdb27eb240L59 which includes an update to change this to a float for the NeighborLoader, can we aim to do the same here once that lands? 
created_at: 2025-05-09T20:14:23Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +282,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: int = 60,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous NABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous NABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: nit: update to ABLP, ditto elsewhere
created_at: 2025-05-09T20:16:15Z
author_association: COLLABORATOR
diff_hunk: @@ -207,6 +255,20 @@ def __init__(
             pin_memory=device.type == "cuda",
         )
 
+        # Eventually - we should expose this as an arg to be passed in.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: done.
created_at: 2025-05-09T21:37:41Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +145,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            transforms (Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]): A sequence of transforms to apply to the data.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: PyG takes in `transform` [1] so I don't think it's too odd to add this? I think that if we want to implement pre-fetching in the future then we'd also want to pre-run the transforms here right? So it makes sense to store them as part of the class.

Is the concern here that it's confusing why the DistABLPLoader is using transforms?

[1]: https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/loader/node_loader.html#NodeLoader
created_at: 2025-05-09T21:40:23Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +282,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: int = 60,
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: updated. 
created_at: 2025-05-09T21:41:05Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +282,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: int = 60,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous NABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous NABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Done.
created_at: 2025-05-09T21:41:34Z
author_association: COLLABORATOR
diff_hunk: @@ -167,6 +167,103 @@ def test_distributed_neighbor_loader_heterogeneous(self):
 
         self.assertEqual(count, 4057)
 
+    @parameterized.expand(
+        [
+            param(
+                "Positive and Negative edges",
+                labeled_edges={
+                    _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]),
+                    _NEGATIVE_EDGE_TYPE: torch.tensor([[10, 11], [16, 14]]),
+                },
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 16, 16, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 12, 14, 13, 17]),
+            ),
+            param(
+                "Positive edges",
+                labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10], [15]])},
+                expected_node=torch.tensor([10, 11, 12, 13, 14, 15, 17]),
+                expected_srcs=torch.tensor([10, 10, 15, 15, 11, 11]),
+                expected_dsts=torch.tensor([11, 12, 13, 14, 13, 17]),
+            ),
+        ]
+    )
+    def test_distributed_neighbor_loader_with_supervision_edges(
path: python/tests/unit/distributed/distributed_neighborloader_test.py
user: login: kmontemayor2-sc
body: sure, done.
created_at: 2025-05-09T21:41:53Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +145,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            transforms (Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]): A sequence of transforms to apply to the data.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: I'd vote for keeping DistNeighborLoader intact from the glt version, since it's going to be used for other cases say random neg, node classification, link prediction, unless we'd expect transforms in those cases. Anything specific to ABLP can go into ABLP loader. 
created_at: 2025-05-09T22:04:15Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +145,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            transforms (Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]): A sequence of transforms to apply to the data.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: Gotcha, makes sense to conform to PyG here. I think my confusion here is that it feels weird that `DistNeighborLoader` is currently accepting `transform` as an argument but we don't have an immediate use case for that class right now, while the `DistABLPLoader`, which we do have a `transform` use-case ready for, doesn't take this as an argument to `__init__`. I feel like if we have `transform` accepted as an argument here for `DistNeighborLoader`, we should be consistent and also have `DistABLPLoader` do the same. 
created_at: 2025-05-09T22:13:03Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: For het cases, node ids would still just be a tensor of anchor nodes (no mapping), since we'd likely have separate loaders for different supervision edge types(that has different anchor nodes). Dataset will be heterogeneous tho. We can already support that in this PR?

But we do need EdgeType mapping for num neighbors. If we can't support it rn it's ok to leave as TODO. 
created_at: 2025-05-09T22:20:30Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: What do we do for homogeneous SSL cases? We'd create the labels and het dataset in partitioner? 
created_at: 2025-05-09T22:23:34Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
+            )
+        if DEFAULT_HOMOGENEOUS_EDGE_TYPE not in dataset.graph:
+            raise ValueError(
+                f"With Homogeneous ABLP, the graph must have {DEFAULT_HOMOGENEOUS_EDGE_TYPE} edge type, received {dataset.graph.keys()}."
+            )
+        positive_label_edge_type, negative_label_edge_type = select_label_edge_types(
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE, dataset.graph.keys()
+        )
+        positive_labels, negative_labels = get_labels_for_anchor_nodes(
+            dataset, input_nodes, positive_label_edge_type, negative_label_edge_type
+        )
+
+        if negative_labels is not None:
+            input_nodes = torch.cat(
+                [input_nodes.unsqueeze(1), positive_labels, negative_labels], dim=1
+            )
+        else:
+            input_nodes = torch.cat([input_nodes.unsqueeze(1), positive_labels], dim=1)
+        input_nodes = (DEFAULT_HOMOGENEOUS_NODE_TYPE, input_nodes)
+        logger.info(
+            f"Converted input nodes to tuple of ({DEFAULT_HOMOGENEOUS_NODE_TYPE}, {input_nodes[1].shape})."
+        )
+        transforms = [
+            _SupervisedToHomogeneous(
+                message_passing_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
+            )
+        ]
+        # TODO(kmonte): stop setting fanout for positive/negative once GLT sampling is fixed.
+        zero_samples = [0] * len(num_neighbors)
path: python/gigl/distributed/distributed_neighborloader.py
user: login: yliu2-sc
body: does this mean that we don't sample for positive/negative edge types? ex. we have a user-story graph. Pos edge type is user-story. For user nodes we want sample user-story-user. Positive story nodes we want to sample story-user-story. How would the zero samples work in this case? 
created_at: 2025-05-09T22:30:29Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
+            )
+        if DEFAULT_HOMOGENEOUS_EDGE_TYPE not in dataset.graph:
+            raise ValueError(
+                f"With Homogeneous ABLP, the graph must have {DEFAULT_HOMOGENEOUS_EDGE_TYPE} edge type, received {dataset.graph.keys()}."
+            )
+        positive_label_edge_type, negative_label_edge_type = select_label_edge_types(
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE, dataset.graph.keys()
+        )
+        positive_labels, negative_labels = get_labels_for_anchor_nodes(
+            dataset, input_nodes, positive_label_edge_type, negative_label_edge_type
+        )
+
+        if negative_labels is not None:
+            input_nodes = torch.cat(
+                [input_nodes.unsqueeze(1), positive_labels, negative_labels], dim=1
+            )
+        else:
+            input_nodes = torch.cat([input_nodes.unsqueeze(1), positive_labels], dim=1)
+        input_nodes = (DEFAULT_HOMOGENEOUS_NODE_TYPE, input_nodes)
+        logger.info(
+            f"Converted input nodes to tuple of ({DEFAULT_HOMOGENEOUS_NODE_TYPE}, {input_nodes[1].shape})."
+        )
+        transforms = [
+            _SupervisedToHomogeneous(
+                message_passing_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
+            )
+        ]
+        # TODO(kmonte): stop setting fanout for positive/negative once GLT sampling is fixed.
+        zero_samples = [0] * len(num_neighbors)
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: My understanding is that this is only setting the `zero_sample` for the `user -> positive_to -> story` relation type. We would still have `user->to->story` and `story->to->user` edges with non-zero fanouts, which would provide the subgraphs for `user` and `story` nodes. This is done so that when we are hopping out from a user, we don't accidentally use the supervision edge, which should only be used as label metadata IIUC.
created_at: 2025-05-09T22:37:06Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: Correct this is done in the `dataset_factory` [here](https://github.com/Snapchat/GiGL/blob/main/python/gigl/distributed/dataset_factory.py#L149-L177) currently, I imagine eventually this logic will move to the transductive splitter
created_at: 2025-05-09T22:38:47Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +145,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            transforms (Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]): A sequence of transforms to apply to the data.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: I feel like "transforms" are pretty generic, no? if we feel strongly I can move all of this to the `DistABLPLoader` but I feel like it's not a big deviation from PyG to have these (or even torch if we thing of `collate_fn` as a "transform".


We actually don't have a usecase for `DistABLPLoader` to take in transform as an argument, just for it to automatically generate one transform if needed.

I prefer to not add `transforms` as an arg to `DistABLPLoader` until we actually have a usecase for user-provided transforms.
created_at: 2025-05-09T23:15:56Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: for het cases users need to provide the node type to iterate over right, like `(UserNodeType, UserNodeIds)`. Otherwise, how do we know which type of anchor nodes to fan out from? Node ids are not unique across node types, we can have `User, 0` and `Item, 0` as valid `NodeType, NodeId` tuples. 


created_at: 2025-05-09T23:17:34Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
+            )
+        if DEFAULT_HOMOGENEOUS_EDGE_TYPE not in dataset.graph:
+            raise ValueError(
+                f"With Homogeneous ABLP, the graph must have {DEFAULT_HOMOGENEOUS_EDGE_TYPE} edge type, received {dataset.graph.keys()}."
+            )
+        positive_label_edge_type, negative_label_edge_type = select_label_edge_types(
+            DEFAULT_HOMOGENEOUS_EDGE_TYPE, dataset.graph.keys()
+        )
+        positive_labels, negative_labels = get_labels_for_anchor_nodes(
+            dataset, input_nodes, positive_label_edge_type, negative_label_edge_type
+        )
+
+        if negative_labels is not None:
+            input_nodes = torch.cat(
+                [input_nodes.unsqueeze(1), positive_labels, negative_labels], dim=1
+            )
+        else:
+            input_nodes = torch.cat([input_nodes.unsqueeze(1), positive_labels], dim=1)
+        input_nodes = (DEFAULT_HOMOGENEOUS_NODE_TYPE, input_nodes)
+        logger.info(
+            f"Converted input nodes to tuple of ({DEFAULT_HOMOGENEOUS_NODE_TYPE}, {input_nodes[1].shape})."
+        )
+        transforms = [
+            _SupervisedToHomogeneous(
+                message_passing_edge_type=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
+            )
+        ]
+        # TODO(kmonte): stop setting fanout for positive/negative once GLT sampling is fixed.
+        zero_samples = [0] * len(num_neighbors)
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: My understanding is the same as Matt's here. 
created_at: 2025-05-09T23:19:25Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +145,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            transforms (Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]): A sequence of transforms to apply to the data.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: mkolodner-sc
body: > We actually don't have a usecase for DistABLPLoader to take in transform as an argument, just for it to automatically generate one transform if needed.

Hmm, I feel like we should be consistent between either both `DistABLPLoader` and `DistNeighborLoader` taking a `transform` argument, or neither take a `transform` argument. If neither `DistABLPLoader` or `DistNeighborLoader` require the user to provide `self._transforms` at this time and GLTNeighborLoader doesn't have this [[1](https://github.com/alibaba/graphlearn-for-pytorch/blob/main/graphlearn_torch/python/distributed/dist_neighbor_loader.py#L75-L87)], can we remove this as an argument altogether?
created_at: 2025-05-09T23:27:11Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +49,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _SupervisedToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: zfan3-sc
body: Do we want to expose this class to the user and let user explicitly provide it? If not how do we know when to treat the input graph as a homogeneous graph with labels v.s. a true heterogeneous graph? 
created_at: 2025-05-10T00:19:48Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +49,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _SupervisedToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Hmmm, for now we are doing this is we see that there's one positive label, and (optionally) a negative label type. 

We can have this be explicitly exposed instead - if you're ok with needing to use this yourself maybe that'd be simpler. 

Ultimately depends how little users want o be exposed to the "heterogeneous" representation I guess :)
created_at: 2025-05-12T16:59:28Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
+            )
+        if DEFAULT_HOMOGENEOUS_EDGE_TYPE not in dataset.graph:
+            raise ValueError(
+                f"With Homogeneous ABLP, the graph must have {DEFAULT_HOMOGENEOUS_EDGE_TYPE} edge type, received {dataset.graph.keys()}."
+            )
+        positive_label_edge_type, negative_label_edge_type = select_label_edge_types(
path: python/gigl/distributed/distributed_neighborloader.py
user: login: zfan3-sc
body: Not necessarily need to be addressed in this PR. Just want to call out ideally we should be able to have multiple positive and negative edge types. 
created_at: 2025-05-12T17:06:51Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +49,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _SupervisedToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: zfan3-sc
body: > we are doing this is we see that there's one positive label, and (optionally) a negative label type.

I see, then how do we deal with heterogeneous graph trained with supervised labels? In that case we have positive label do not need this transformation
created_at: 2025-05-12T17:07:03Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +49,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _SupervisedToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: For a truly heterogeneous graph (containing e.g. User and Content node types), then we don't want to use this right?

My understanding is that this should *only* be used for "psuedo-heterogeneous" graphs where there's only one node type, and one message passing edge type. 

For cases where we have multiple sets of label edge types I think we can extend this.

If there are multiple "true" node or edge types in what case would we want to convert to homogeneous Data?. 
created_at: 2025-05-12T17:13:33Z
author_association: COLLABORATOR
diff_hunk: @@ -109,6 +145,7 @@ def __init__(
                 Defaults to `2` if set to `None` when using cpu training/inference.
             shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
             drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+            transforms (Sequence[Callable[[Union[Data, HeteroData]], Union[Data, HeteroData]]]): A sequence of transforms to apply to the data.
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: Sure, I guess we can remove for now (though I feel like it makes our code a bit more complicated). 
created_at: 2025-05-12T17:14:02Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +281,175 @@ def __init__(
             edge_dir=dataset.edge_dir,
             seed=None,  # it's actually optional - None means random.
         )
-        if isinstance(curr_process_nodes, torch.Tensor):
-            node_ids = curr_process_nodes
-            node_type = None
-        else:
-            node_type, node_ids = curr_process_nodes
+        super().__init__(dataset, input_data, sampling_config, device, worker_options)
 
-        input_data = _BatchedNodeSamplerInput(node=node_ids, input_type=node_type)
+    def _collate_fn(self, msg: SampleMessage) -> Union[Data, HeteroData]:
+        data = super()._collate_fn(msg)
+        for transform in self._transforms:
+            data = transform(data)
+        return data
 
-        super().__init__(dataset, input_data, sampling_config, device, worker_options)
+
+class DistABLPLoader(DistNeighborLoader):
+    def __init__(
+        self,
+        dataset: DistLinkPredictionDataset,
+        num_neighbors: Union[List[int], Dict[EdgeType, List[int]]],
+        context: DistributedContext,
+        local_process_rank: int,  # TODO: Move this to DistributedContext
+        local_process_world_size: int,  # TODO: Move this to DistributedContext
+        input_nodes: Optional[
+            Union[torch.Tensor, Tuple[NodeType, torch.Tensor]]
+        ] = None,
+        num_workers: int = 1,
+        batch_size: int = 1,
+        pin_memory_device: Optional[torch.device] = None,
+        worker_concurrency: int = 4,
+        channel_size: str = "4GB",
+        process_start_gap_seconds: float = 60.0,
+        num_cpu_threads: Optional[int] = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ):
+        """
+        Neighbor loader for Anchor Based Link Prediction (ABLP) tasks.
+
+        Note that for this class, the dataset must *always* be heterogeneous,
+        as we need separate edge types for positive and negative labels.
+
+        If you provide `input_nodes` for homogeneous input (only as a Tensor),
+        Then we will attempt to infer the positive and optional negative labels
+        from the dataset.
+        In this case, the output of the loader will be a torch_geometric.data.Data object.
+        Otherwise, the output will be a torch_geometric.data.HeteroData object.
+
+           Args:
+            dataset (DistLinkPredictionDataset): The dataset to sample from.
+            num_neighbors (List[int] or Dict[Tuple[str, str, str], List[int]]):
+                The number of neighbors to sample for each node in each iteration.
+                If an entry is set to `-1`, all neighbors will be included.
+                In heterogeneous graphs, may also take in a dictionary denoting
+                the amount of neighbors to sample for each individual edge type.
+            context (DistributedContext): Distributed context information of the current process.
+            local_process_rank (int): The local rank of the current process within a node.
+            local_process_world_size (int): The total number of processes within a node.
+            input_nodes (torch.Tensor or Tuple[str, torch.Tensor]): The
+                indices of seed nodes to start sampling from.
+                It is of type `torch.LongTensor` for homogeneous graphs.
+                If set to `None` for homogeneous settings, all nodes will be considered.
+                In heterogeneous graphs, this flag must be passed in as a tuple that holds
+                the node type and node indices. (default: `None`)
+            num_workers (int): How many workers to use (subprocesses to spwan) for
+                    distributed neighbor sampling of the current process. (default: ``1``).
+            batch_size (int, optional): how many samples per batch to load
+                (default: ``1``).
+            pin_memory_device (str, optional): The target device that the sampled
+                results should be copied to. If set to ``None``, the device is inferred based off of
+                (got by ``gigl.distributed.utils.device.get_available_device``). Which uses the
+                local_process_rank and torch.cuda.device_count() to assign the device. If cuda is not available,
+                the cpu device will be used. (default: ``None``).
+            worker_concurrency (int): The max sampling concurrency for each sampling
+                worker. Load testing has showed that setting worker_concurrency to 4 yields the best performance
+                for sampling. Although, you may whish to explore higher/lower settings when performance tuning.
+                (default: `4`).
+            channel_size (int or str): The shared-memory buffer size (bytes) allocated
+                for the channel. Can be modified for performance tuning; a good starting point is: ``num_workers * 64MB``
+                (default: "4GB").
+            process_start_gap_seconds (float): Delay between each process for initializing neighbor loader. At large scales,
+                it is recommended to set this value to be between 60 and 120 seconds -- otherwise multiple processes may
+                attempt to initialize dataloaders at overlapping times, which can cause CPU memory OOM.
+            num_cpu_threads (Optional[int]): Number of cpu threads PyTorch should use for CPU training/inference
+                neighbor loading; on top of the per process parallelism.
+                Defaults to `2` if set to `None` when using cpu training/inference.
+            shuffle (bool): Whether to shuffle the input nodes. (default: ``False``).
+            drop_last (bool): Whether to drop the last incomplete batch. (default: ``False``).
+        """
+
+        # Set self._shutdowned right away, that way if we throw here, and __del__ is called,
+        # then we can properly clean up and don't get extraneous error messages.
+        # We set to `True` as we don't need to cleanup right away, and this will get set
+        # to `False` in super().__init__()` e.g.
+        # https://github.com/alibaba/graphlearn-for-pytorch/blob/26fe3d4e050b081bc51a79dc9547f244f5d314da/graphlearn_torch/python/distributed/dist_loader.py#L125C1-L126C1
+        self._shutdowned = True
+        # TODO(kmonte): Remove these checks and support properly heterogeneous NABLP.
+        if isinstance(input_nodes, tuple):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide node_ids as a tensor."
+            )
+        if isinstance(num_neighbors, abc.Mapping):
+            raise ValueError(
+                "Heterogeneous ABLP is not supported yet. Provide num_neighbors as a list of integers."
+            )
+        if input_nodes is None:
+            if dataset.node_ids is None:
+                raise ValueError(
+                    "Dataset must have node ids if input_nodes are not provided."
+                )
+            if isinstance(dataset.node_ids, abc.Mapping):
+                raise ValueError(
+                    f"input_nodes must be provided for heterogeneous datasets, received node_ids of type: {dataset.node_ids.keys()}"
+                )
+            input_nodes = dataset.node_ids
+        if len(input_nodes.shape) != 1:
+            raise ValueError(
+                f"input_nodes must be a 1D tensor, got {input_nodes.shape}."
+            )
+        if not isinstance(dataset.graph, abc.Mapping):
+            raise ValueError(
+                f"The dataset must be heterogeneous for ABLP. Recieved dataset with graph of type: {type(dataset.graph)}"
+            )
+        if DEFAULT_HOMOGENEOUS_EDGE_TYPE not in dataset.graph:
+            raise ValueError(
+                f"With Homogeneous ABLP, the graph must have {DEFAULT_HOMOGENEOUS_EDGE_TYPE} edge type, received {dataset.graph.keys()}."
+            )
+        positive_label_edge_type, negative_label_edge_type = select_label_edge_types(
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: ack yes, we can do so as a follow up :)
created_at: 2025-05-12T17:18:51Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +49,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _SupervisedToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: zfan3-sc
body: > If there are multiple "true" node or edge types in what case would we want to convert to homogeneous Data?

Ah no. My understanding was this dataloader is gonna be used by both "true" het graph and homogeneous graph with supervision labels. If this is only for latter I think it's fine as is. 
created_at: 2025-05-12T17:32:16Z
author_association: COLLABORATOR
diff_hunk: @@ -40,6 +49,30 @@ def __getitem__(
         return _BatchedNodeSamplerInput(self.node[index].view(-1), self.input_type)
 
 
+class _SupervisedToHomogeneous:
path: python/gigl/distributed/distributed_neighborloader.py
user: login: kmontemayor2-sc
body: We do intend to use this dataloader for "true" het graphs - just that this transform shouldn't be applied to "true" graphs
created_at: 2025-05-12T18:15:26Z
author_association: COLLABORATOR

PR comments:
user: login: kmontemayor2-sc
created_at: 2025-05-02T20:41:29Z
body: /unit_test
author_association: COLLABORATOR
user: login: kmontemayor2-sc
created_at: 2025-05-02T20:41:35Z
body: /integration_test
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-05-02T20:41:44Z
body: # GiGL Automation

@ _20:41:44UTC_ : 🔄 `Unit Test` [started](https://github.com/Snapchat/GiGL/actions/runs/14803036290).

@ _23:43:15UTC_ : ❌ Workflow failed.
Please check the logs for more details.

author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-02T20:41:47Z
body: # GiGL Automation

@ _20:41:47UTC_ : 🔄 `Integration Test` [started](https://github.com/Snapchat/GiGL/actions/runs/14803037824).

@ _21:33:11UTC_ : ✅ Workflow completed successfully.

author_association: CONTRIBUTOR
user: login: kmontemayor2-sc
created_at: 2025-05-09T21:43:07Z
body: /unit_test
author_association: COLLABORATOR
user: login: kmontemayor2-sc
created_at: 2025-05-09T21:43:13Z
body: /integration_test
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-05-09T21:43:23Z
body: # GiGL Automation

@ _21:43:23UTC_ : 🔄 `Unit Test` [started](https://github.com/Snapchat/GiGL/actions/runs/14938230595).

@ _24:45:02UTC_ : ❌ Workflow failed.
Please check the logs for more details.

author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-09T21:43:26Z
body: # GiGL Automation

@ _21:43:25UTC_ : 🔄 `Integration Test` [started](https://github.com/Snapchat/GiGL/actions/runs/14938231426).

@ _22:39:25UTC_ : ✅ Workflow completed successfully.

author_association: CONTRIBUTOR
user: login: kmontemayor2-sc
created_at: 2025-05-12T16:12:38Z
body: /unit_test
author_association: COLLABORATOR
user: login: kmontemayor2-sc
created_at: 2025-05-12T16:12:46Z
body: /integration_test
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-05-12T16:12:52Z
body: # GiGL Automation

@ _16:12:52UTC_ : 🔄 `Unit Test` [started](https://github.com/Snapchat/GiGL/actions/runs/14977220639).

@ _16:32:15UTC_ : ✅ Workflow completed successfully.

author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-12T16:13:01Z
body: # GiGL Automation

@ _16:13:01UTC_ : 🔄 `Integration Test` [started](https://github.com/Snapchat/GiGL/actions/runs/14977223341).

@ _16:54:53UTC_ : ✅ Workflow completed successfully.

author_association: CONTRIBUTOR