diff --git a/examples/link_prediction/graph_store/heterogeneous_training.py b/examples/link_prediction/graph_store/heterogeneous_training.py index 2be34e608..2c772236b 100644 --- a/examples/link_prediction/graph_store/heterogeneous_training.py +++ b/examples/link_prediction/graph_store/heterogeneous_training.py @@ -211,6 +211,7 @@ def _setup_dataloaders( channel_size=sampling_worker_shared_channel_size, process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) print(f"---Rank {rank} finished setting up main loader for split={split}") @@ -294,23 +295,16 @@ def _compute_loss( ) # Extracting local query, random negative, positive, hard_negative, and random_negative indices. - query_node_idx: torch.Tensor = torch.arange( - main_data[query_node_type].batch_size - ).to(device) + query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device) random_negative_batch_size = random_negative_data[labeled_node_type].batch_size - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/graph_store/homogeneous_training.py b/examples/link_prediction/graph_store/homogeneous_training.py index 2d4c22788..25ece94a0 100644 --- a/examples/link_prediction/graph_store/homogeneous_training.py +++ b/examples/link_prediction/graph_store/homogeneous_training.py @@ -241,6 +241,7 @@ def _setup_dataloaders( channel_size=sampling_worker_shared_channel_size, process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader for split={split}") @@ -302,21 +303,16 @@ def _compute_loss( main_embeddings = model(data=main_data, device=device) random_negative_embeddings = model(data=random_negative_data, device=device) - query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device) + query_node_idx = torch.arange(main_data.batch_size, device=device) random_negative_batch_size = random_negative_data.batch_size - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/heterogeneous_training.py b/examples/link_prediction/heterogeneous_training.py index 8ed672b7c..e41a19191 100644 --- a/examples/link_prediction/heterogeneous_training.py +++ b/examples/link_prediction/heterogeneous_training.py @@ -144,6 +144,7 @@ def _setup_dataloaders( # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader") @@ -218,25 +219,16 @@ def _compute_loss( # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset. # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique. # TODO (mkolodner-sc): If GLT migrates towards using PyG's `.n_id` field instead of `.node`, update this code. - query_node_idx: torch.Tensor = torch.arange( - main_data[query_node_type].batch_size - ).to(device) + query_node_idx = torch.arange(main_data[query_node_type].batch_size, device=device) random_negative_batch_size = random_negative_data[labeled_node_type].batch_size - # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor], even in the heterogeneous setting. - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/examples/link_prediction/homogeneous_training.py b/examples/link_prediction/homogeneous_training.py index b95a77489..a84b55114 100644 --- a/examples/link_prediction/homogeneous_training.py +++ b/examples/link_prediction/homogeneous_training.py @@ -134,6 +134,7 @@ def _setup_dataloaders( # This is done so that each process on the current machine which initializes a `main_loader` doesn't compete for memory, causing potential OOM process_start_gap_seconds=process_start_gap_seconds, shuffle=shuffle, + use_edge_index_output=True, ) logger.info(f"---Rank {rank} finished setting up main loader") @@ -187,23 +188,17 @@ def _compute_loss( # Extracting local query, random negative, positive, hard_negative, and random_negative indices. # Local in this case refers to the local index in the batch, while global subsequently refers to the node's unique global ID across all nodes in the dataset. # Global ids are stored in data.node, ex. `data.node = [50, 20, 10]` where the `0` is the local index for global id `50`. Note that each global id in data.node is unique. - query_node_idx: torch.Tensor = torch.arange(main_data.batch_size).to(device) + query_node_idx = torch.arange(main_data.batch_size, device=device) random_negative_batch_size = random_negative_data.batch_size - # main_data.y_positive is a dict[query_node_local_index: int, labeled_node_local_indices: torch.Tensor] - positive_idx: torch.Tensor = torch.cat(list(main_data.y_positive.values())).to( - device - ) - # We also extract a repeated query node index tensor which upsamples each query node based on the number of positives it has - repeated_query_node_idx = query_node_idx.repeat_interleave( - torch.tensor([len(v) for v in main_data.y_positive.values()]).to(device) - ) + # Label edge indices are [anchor_local_id, label_local_id] pairs. + positive_label_edge_index: torch.Tensor = main_data.y_positive + repeated_query_node_idx = query_node_idx[positive_label_edge_index[0]] + positive_idx = positive_label_edge_index[1] if hasattr(main_data, "y_negative"): - hard_negative_idx: torch.Tensor = torch.cat( - list(main_data.y_negative.values()) - ).to(device) + hard_negative_idx: torch.Tensor = main_data.y_negative[1] else: - hard_negative_idx = torch.empty(0, dtype=torch.long).to(device) + hard_negative_idx = torch.empty(0, dtype=torch.long, device=device) # Use local IDs to get the corresponding embeddings in the tensors diff --git a/gigl/distributed/dist_ablp_neighborloader.py b/gigl/distributed/dist_ablp_neighborloader.py index 50f42f5a9..9dc5e745d 100644 --- a/gigl/distributed/dist_ablp_neighborloader.py +++ b/gigl/distributed/dist_ablp_neighborloader.py @@ -1,4 +1,5 @@ -from collections import abc, defaultdict +import warnings +from collections import abc from itertools import count from typing import Optional, Union @@ -27,6 +28,10 @@ SamplerOptions, resolve_sampler_options, ) +from gigl.distributed.utils.ablp import ( + label_edge_index_to_dict, + remap_labels_to_local_edge_indices, +) from gigl.distributed.utils.neighborloader import ( DatasetSchema, SamplingClusterSetup, @@ -43,7 +48,6 @@ from gigl.types.graph import ( DEFAULT_HOMOGENEOUS_EDGE_TYPE, DEFAULT_HOMOGENEOUS_NODE_TYPE, - label_edge_type_to_message_passing_edge_type, message_passing_to_negative_label, message_passing_to_positive_label, reverse_edge_type, @@ -90,17 +94,19 @@ def __init__( local_process_rank: Optional[int] = None, # TODO: (svij) Deprecate this local_process_world_size: Optional[int] = None, # TODO: (svij) Deprecate this non_blocking_transfers: bool = True, + use_edge_index_output: 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. + The dataset must *always* be heterogeneous here, since positive and + negative labels are carried as separate edge types. By default, the loader will return {py:class} `torch_geometric.data.HeteroData` (heterogeneous) objects, but will return a {py:class}`torch_geometric.data.Data` (homogeneous) object if the dataset is "labeled homogeneous". - The following fields may also be present: + The following fields may also be present (this describes the deprecated + default shape; see ``use_edge_index_output`` for the tensor format): - `y_positive`: `dict[int, torch.Tensor]` mapping from local anchor node id to a tensor of positive label node ids. - `y_negative`: (Optional) `dict[int, torch.Tensor]` mapping from local anchor node id to a tensor of negative @@ -138,6 +144,16 @@ def __init__( - `y_positive`: {(a, to, b): {0: torch.tensor([1])}, (a, to, c): {0: torch.tensor([2])}} - `y_negative`: {(a, to, b): {0: torch.tensor([3])}, (a, to, c): {0: torch.tensor([4])}} + With ``use_edge_index_output=True``, each label field is a ``[2, E]`` + tensor. Row 0 contains local anchor indices and row 1 contains local + label-node indices. For the example above: + + - ``y_positive = torch.tensor([[0], [1]])`` + - ``y_negative = torch.tensor([[0], [2]])`` + + Multiple supervision edge types produce ``dict[EdgeType, torch.Tensor]``. + Pair order within an anchor is unspecified. + Args: dataset (Union[DistDataset, RemoteDistDataset]): The dataset to sample from. If this is a `RemoteDistDataset`, then we are in "Graph Store" mode. @@ -213,11 +229,23 @@ def __init__( is used instead. See https://docs.pytorch.org/tutorials/intermediate/pinmem_nonblock.html for background on pinned memory and non-blocking transfers. + use_edge_index_output (bool): Return labels as ``[2, E]`` edge-index + tensors instead of the deprecated ragged dictionaries. Row 0 + contains local anchor indices and row 1 contains local label-node + indices. Defaults to ``False`` for backward compatibility. """ # 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. self._shutdowned = True + if not use_edge_index_output: + warnings.warn( + "The ragged dictionary ABLP label output is deprecated. Pass " + "use_edge_index_output=True to receive [2, E] label edge indices.", + FutureWarning, + stacklevel=2, + ) + self._use_edge_index_output = use_edge_index_output sampler_options = resolve_sampler_options(num_neighbors, sampler_options) @@ -752,71 +780,99 @@ def _set_labels( positive_labels_by_label_edge_type: dict[EdgeType, torch.Tensor], negative_labels_by_label_edge_type: dict[EdgeType, torch.Tensor], ) -> Union[Data, HeteroData]: - """ - Sets the labels and relevant fields in the torch_geometric Data object, converting the global node ids for labels to their - local index. Removes inserted supervision edge type from the data variables, since this is an implementation detail and should not be - exposed in the final HeteroData/Data object. + """Attach ABLP labels to the collated graph, remapped to subgraph-local indices. + + The tensor output uses ``[2, E]`` label edge indices. The compatibility + path converts the same tensors to the deprecated ragged dictionaries. + Args: - data (Union[Data, HeteroData]): Graph to provide labels for - positive_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Dict[positive label edge type, label ID tensor], - where the ith row of the tensor corresponds to the ith anchor node ID. - negative_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Dict[negative label edge type, label ID tensor], - where the ith row of the tensor corresponds to the ith anchor node ID. + data (Union[Data, HeteroData]): Graph to attach labels to. + positive_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): Per + positive-label edge type, a ``[N_anchors, M]`` tensor whose ``i``-th + row holds the global label ids of the ``i``-th anchor. + negative_labels_by_label_edge_type (dict[EdgeType, torch.Tensor]): As + above, for negative-label edge types. + Returns: - Union[Data, HeteroData]: torch_geometric HeteroData/Data object with the filtered edge fields and labels set as properties of the instance + Union[Data, HeteroData]: The same object with the supervision edge fields + stripped and ``y_positive`` (and ``y_negative`` when present) attached. + + Raises: + ValueError: If no positive labels are found in ``data``. """ - # shape [N], where N is the number of nodes in the subgraph, and local_node_to_global_node[i] gives the global node id for local node id `i` - node_type_to_local_node_to_global_node: dict[NodeType, torch.Tensor] = {} + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor] = {} if isinstance(data, HeteroData): - for e_type in self._supervision_edge_types: - node_type_to_local_node_to_global_node[e_type[0]] = data[e_type[0]].node - node_type_to_local_node_to_global_node[e_type[2]] = data[e_type[2]].node + # Key the map off the label edge types rather than + # self._supervision_edge_types: the latter is stored reversed when + # edge_dir="in" (see __init__), so its dst node type is the anchor + # type there, while label edge types always arrive in outward form. + # Deriving from the keys that _remap_labels_by_edge_type will look up + # keeps the two in step regardless of edge direction. + for label_edge_type, label_tensor in ( + *positive_labels_by_label_edge_type.items(), + *negative_labels_by_label_edge_type.items(), + ): + # Mirror the zero-anchor skip in _remap_labels_by_edge_type. An + # absent node type would otherwise be auto-created as an empty + # store by HeteroData.__getitem__, turning a clear KeyError into + # an AttributeError on .node. + if label_tensor.size(0) == 0: + continue + supervision_node_type = label_edge_type[2] + local_id_to_global_id_by_node_type[supervision_node_type] = data[ + supervision_node_type + ].node else: - node_type_to_local_node_to_global_node[DEFAULT_HOMOGENEOUS_NODE_TYPE] = ( + local_id_to_global_id_by_node_type[DEFAULT_HOMOGENEOUS_NODE_TYPE] = ( data.node ) - output_positive_labels: dict[EdgeType, dict[int, torch.Tensor]] = defaultdict( - dict - ) - output_negative_labels: dict[EdgeType, dict[int, torch.Tensor]] = defaultdict( - dict + + ( + positive_label_edge_indices, + negative_label_edge_indices, + ) = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + positive_labels_by_edge_type=positive_labels_by_label_edge_type, + negative_labels_by_edge_type=negative_labels_by_label_edge_type, ) - # We always have supervision edge types of the form (anchor_node_type, to, supervision_node_type) - # So we can index into the edge type accordingly. - edge_index = 2 - for edge_type, label_tensor in positive_labels_by_label_edge_type.items(): - for local_anchor_node_id in range(label_tensor.size(0)): - positive_mask = ( - node_type_to_local_node_to_global_node[ - edge_type[edge_index] - ].unsqueeze(1) - == label_tensor[local_anchor_node_id] - ) # shape [N, P], where N is the number of nodes and P is the number of positive labels for the current anchor node - - # Gets the indexes of the items in local_node_to_global_node which match any of the positive labels for the current anchor node - output_positive_labels[ - label_edge_type_to_message_passing_edge_type(edge_type) - ][local_anchor_node_id] = torch.nonzero(positive_mask)[:, 0].to( - self.to_device + output_positive_labels: dict[ + EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]] + ] = {} + output_negative_labels: dict[ + EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]] + ] = {} + if self._use_edge_index_output: + for edge_type, label_edge_index in positive_label_edge_indices.items(): + output_positive_labels[edge_type] = label_edge_index + for edge_type, label_edge_index in negative_label_edge_indices.items(): + output_negative_labels[edge_type] = label_edge_index + else: + label_edge_types = ( + positive_label_edge_indices.keys() | negative_label_edge_indices.keys() + ) + if isinstance(data, HeteroData): + num_anchors_by_edge_type = { + edge_type: int(data[edge_type[0]].batch_size) + for edge_type in label_edge_types + } + else: + num_anchors_by_edge_type = { + edge_type: int(data.batch_size) for edge_type in label_edge_types + } + output_positive_labels = { + edge_type: label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=num_anchors_by_edge_type[edge_type], ) - # Shape [X], where X is the number of indexes in the original local_node_to_global_node which match a node in the positive labels for the current anchor node - - for edge_type, label_tensor in negative_labels_by_label_edge_type.items(): - for local_anchor_node_id in range(label_tensor.size(0)): - negative_mask = ( - node_type_to_local_node_to_global_node[ - edge_type[edge_index] - ].unsqueeze(1) - == label_tensor[local_anchor_node_id] - ) # shape [N, M], where N is the number of nodes and M is the number of negative labels for the current anchor node - - # Gets the indexes of the items in local_node_to_global_node which match any of the negative labels for the current anchor node - output_negative_labels[ - label_edge_type_to_message_passing_edge_type(edge_type) - ][local_anchor_node_id] = torch.nonzero(negative_mask)[:, 0].to( - self.to_device + for edge_type, label_edge_index in positive_label_edge_indices.items() + } + output_negative_labels = { + edge_type: label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=num_anchors_by_edge_type[edge_type], ) - # Shape [X], where X is the number of indexes in the original local_node_to_global_node which match a node in the negative labels for the current anchor node + for edge_type, label_edge_index in negative_label_edge_indices.items() + } if not output_positive_labels: raise ValueError("No positive labels were found in the data!") elif len(output_positive_labels) == 1: diff --git a/gigl/distributed/utils/ablp.py b/gigl/distributed/utils/ablp.py new file mode 100644 index 000000000..aeff1e569 --- /dev/null +++ b/gigl/distributed/utils/ablp.py @@ -0,0 +1,160 @@ +"""Utilities for remapping ABLP labels to sampled-subgraph indices.""" + +import torch +from torch_geometric.typing import EdgeType + +from gigl.src.common.types.graph_data import NodeType +from gigl.types.graph import label_edge_type_to_message_passing_edge_type +from gigl.utils.data_splitters import PADDING_NODE + + +def label_edge_index_to_dict( + label_edge_index: torch.Tensor, num_anchors: int +) -> dict[int, torch.Tensor]: + """Convert a label edge index to the deprecated per-anchor dictionary. + + The edge index must be grouped by anchor, as guaranteed by + :func:`remap_labels_to_local_edge_indices`. + + Args: + label_edge_index: A ``[2, E]`` tensor. Row 0 contains local anchor + indices and row 1 contains local label-node indices. + num_anchors: Number of anchors in the sampled batch, including anchors + with no labels. + + Returns: + A mapping from every local anchor index to its local label-node indices. + """ + counts = torch.bincount(label_edge_index[0], minlength=num_anchors) + labels_by_anchor = torch.split(label_edge_index[1], counts.tolist()) + return {anchor: labels_by_anchor[anchor] for anchor in range(num_anchors)} + + +def _remap_label_tensor_to_local_edge_index( + label_tensor: torch.Tensor, + sorted_global_node_ids: torch.Tensor, + local_ids_by_sorted_global_id: torch.Tensor, +) -> torch.Tensor: + """Remap one padded global-label tensor to a local label edge index. + + Args: + label_tensor: A ``[N_anchors, M]`` tensor of global label-node ids, + padded with ``PADDING_NODE``. + sorted_global_node_ids: Sorted global ids for the sampled supervision + nodes. + local_ids_by_sorted_global_id: Local node indices in the order given by + ``sorted_global_node_ids``. + + Returns: + A ``[2, E]`` long tensor on the input device. Row 0 contains local + anchor indices and row 1 contains local label-node indices. Labels not + present in the sampled subgraph are omitted. + + Raises: + ValueError: If the sampled node map contains duplicate global ids. + """ + num_anchors = int(label_tensor.size(0)) + num_nodes = int(sorted_global_node_ids.size(0)) + empty_edge_index = torch.empty((2, 0), dtype=torch.long, device=label_tensor.device) + if num_anchors == 0: + return empty_edge_index + + num_labels = int(label_tensor.size(1)) + candidate_global_label_ids = label_tensor.reshape(-1) + candidate_anchor_indices = torch.arange( + num_anchors, device=label_tensor.device + ).repeat_interleave(num_labels) + + is_not_padding = candidate_global_label_ids != PADDING_NODE + candidate_global_label_ids = candidate_global_label_ids[is_not_padding] + candidate_anchor_indices = candidate_anchor_indices[is_not_padding] + if num_nodes == 0 or candidate_global_label_ids.numel() == 0: + return empty_edge_index + + if not bool((sorted_global_node_ids[1:] > sorted_global_node_ids[:-1]).all()): + raise ValueError( + "Vectorized label remapping requires unique global ids in the " + "sampled node map." + ) + + insertion_indices = torch.searchsorted( + sorted_global_node_ids, candidate_global_label_ids + ) + # An absent label larger than every sampled id inserts at num_nodes. Clamp + # before gathering; the exact-match mask below still discards that label. + insertion_indices = insertion_indices.clamp_(max=num_nodes - 1) + is_in_sampled_subgraph = ( + sorted_global_node_ids[insertion_indices] == candidate_global_label_ids + ) + + local_label_indices = local_ids_by_sorted_global_id[insertion_indices][ + is_in_sampled_subgraph + ] + anchor_indices = candidate_anchor_indices[is_in_sampled_subgraph] + return torch.stack((anchor_indices, local_label_indices)) + + +def _remap_labels_by_edge_type( + labels_by_edge_type: dict[EdgeType, torch.Tensor], + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + sorted_node_lookup_by_type: dict[NodeType, tuple[torch.Tensor, torch.Tensor]], +) -> dict[EdgeType, torch.Tensor]: + """Remap positive or negative labels for each supervision edge type.""" + output: dict[EdgeType, torch.Tensor] = {} + for label_edge_type, label_tensor in labels_by_edge_type.items(): + if label_tensor.size(0) == 0: + continue + + supervision_node_type = label_edge_type[2] + if supervision_node_type not in sorted_node_lookup_by_type: + sorted_node_lookup_by_type[supervision_node_type] = torch.sort( + local_id_to_global_id_by_node_type[supervision_node_type] + ) + ( + sorted_global_node_ids, + local_ids_by_sorted_global_id, + ) = sorted_node_lookup_by_type[supervision_node_type] + output[label_edge_type_to_message_passing_edge_type(label_edge_type)] = ( + _remap_label_tensor_to_local_edge_index( + label_tensor=label_tensor, + sorted_global_node_ids=sorted_global_node_ids, + local_ids_by_sorted_global_id=local_ids_by_sorted_global_id, + ) + ) + return output + + +def remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + positive_labels_by_edge_type: dict[EdgeType, torch.Tensor], + negative_labels_by_edge_type: dict[EdgeType, torch.Tensor], +) -> tuple[dict[EdgeType, torch.Tensor], dict[EdgeType, torch.Tensor]]: + """Remap padded global ABLP labels to local label edge indices. + + Positive and negative labels share one sorted node lookup per supervision + node type. Pair order within an anchor is unspecified. + + Args: + local_id_to_global_id_by_node_type: Per node type, a tensor whose + ``i``-th entry is the global id for local node ``i``. + positive_labels_by_edge_type: Per positive-label edge type, a padded + ``[N_anchors, M]`` tensor of global label-node ids. + negative_labels_by_edge_type: Equivalent negative-label tensors. May + be empty. + + Returns: + Positive and negative mappings keyed by message-passing edge type. Each + value is a ``[2, E]`` local label edge index. + """ + sorted_node_lookup_by_type: dict[NodeType, tuple[torch.Tensor, torch.Tensor]] = {} + positive_label_edge_indices = _remap_labels_by_edge_type( + labels_by_edge_type=positive_labels_by_edge_type, + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + sorted_node_lookup_by_type=sorted_node_lookup_by_type, + ) + negative_label_edge_indices = _remap_labels_by_edge_type( + labels_by_edge_type=negative_labels_by_edge_type, + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + sorted_node_lookup_by_type=sorted_node_lookup_by_type, + ) + return positive_label_edge_indices, negative_label_edge_indices diff --git a/tests/unit/distributed/dist_ablp_neighborloader_test.py b/tests/unit/distributed/dist_ablp_neighborloader_test.py index 31d3d1cbc..b7dc1cad4 100644 --- a/tests/unit/distributed/dist_ablp_neighborloader_test.py +++ b/tests/unit/distributed/dist_ablp_neighborloader_test.py @@ -416,6 +416,177 @@ def _run_distributed_ablp_neighbor_loader_multiple_supervision_edge_types( shutdown_rpc() +def _global_pair_set( + anchor_node: torch.Tensor, + label_node: torch.Tensor, + label_dict: dict[int, torch.Tensor], +) -> list[tuple[int, int]]: + """Convert a label dictionary to sorted global-id pairs. + + Anchors and labels are looked up in separate node maps so this works for + heterogeneous graphs, where the two live in different node stores. Pass the + same tensor twice for a homogeneous graph. + """ + pairs: list[tuple[int, int]] = [] + for local_anchor, local_labels in label_dict.items(): + global_anchor = int(anchor_node[local_anchor].item()) + for local_label in local_labels.tolist(): + pairs.append((global_anchor, int(label_node[local_label].item()))) + return sorted(pairs) + + +def _global_pair_set_from_edge_index( + anchor_node: torch.Tensor, + label_node: torch.Tensor, + label_edge_index: torch.Tensor, +) -> list[tuple[int, int]]: + """Convert a label edge index to sorted global-id pairs. + + Takes separate anchor and label node maps for the same reason as + :func:`_global_pair_set`. + """ + return sorted( + ( + int(anchor_node[local_anchor].item()), + int(label_node[local_label].item()), + ) + for local_anchor, local_label in label_edge_index.t().tolist() + ) + + +def _collect_homogeneous_labels( + _, + return_dict, + use_edge_index_output: bool, + dataset: DistDataset, + input_nodes: torch.Tensor, + batch_size: int, + has_negatives: bool, +): + """Run one loader format and return its sorted global-id label pairs.""" + create_test_process_group() + loader = DistABLPLoader( + dataset=dataset, + num_neighbors=[2, 2], + input_nodes=input_nodes, + batch_size=batch_size, + pin_memory_device=torch.device("cpu"), + use_edge_index_output=use_edge_index_output, + ) + positive_pairs: list[tuple[int, int]] = [] + negative_pairs: list[tuple[int, int]] = [] + for datum in loader: + assert isinstance(datum, Data) + node = datum.node + if use_edge_index_output: + assert isinstance(datum.y_positive, torch.Tensor) + assert datum.y_positive.size(0) == 2 + positive_pairs.extend( + _global_pair_set_from_edge_index(node, node, datum.y_positive) + ) + else: + positive_pairs.extend(_global_pair_set(node, node, datum.y_positive)) + if has_negatives: + if use_edge_index_output: + assert isinstance(datum.y_negative, torch.Tensor) + assert datum.y_negative.size(0) == 2 + negative_pairs.extend( + _global_pair_set_from_edge_index(node, node, datum.y_negative) + ) + else: + negative_pairs.extend(_global_pair_set(node, node, datum.y_negative)) + else: + assert not hasattr(datum, "y_negative"), ( + f"expected no negatives, got {getattr(datum, 'y_negative', None)}" + ) + return_dict[use_edge_index_output] = ( + sorted(positive_pairs), + sorted(negative_pairs), + ) + shutdown_rpc() + + +def _edge_type_key(edge_type: EdgeType) -> tuple[str, ...]: + """Canonical dict key for an edge type. + + Label edge types reach collation as plain tuples on some paths and as + ``EdgeType`` on others, and the two stringify differently (``"('a', 'to', + 'b')"`` vs ``'a-to-b'``). Normalizing to plain strings keeps keys comparable + across the process boundary regardless of which arrives. + """ + return tuple(str(part) for part in edge_type) + + +def _accumulate_heterogeneous_pairs( + data: HeteroData, + labels_by_edge_type: dict[EdgeType, Union[torch.Tensor, dict[int, torch.Tensor]]], + use_edge_index_output: bool, + into: dict[tuple[str, ...], list[tuple[int, int]]], +) -> None: + """Accumulate one batch's labels as global (anchor, label) id pairs. + + Anchors and supervision nodes are resolved through the node maps of the edge + type's src and dst stores respectively, so a label remapped against the wrong + node type surfaces as a wrong global id rather than passing silently. + """ + anchor_index = 0 + supervision_index = 2 + for edge_type, labels in labels_by_edge_type.items(): + anchor_node = data[edge_type[anchor_index]].node + label_node = data[edge_type[supervision_index]].node + if use_edge_index_output: + assert isinstance(labels, torch.Tensor), f"{edge_type}: {type(labels)}" + assert labels.size(0) == 2 + pairs = _global_pair_set_from_edge_index(anchor_node, label_node, labels) + else: + assert not isinstance(labels, torch.Tensor), f"{edge_type}: {type(labels)}" + pairs = _global_pair_set(anchor_node, label_node, labels) + into[_edge_type_key(edge_type)].extend(pairs) + + +def _collect_heterogeneous_labels( + _, + return_dict, + use_edge_index_output: bool, + dataset: DistDataset, + input_nodes: tuple[NodeType, torch.Tensor], + supervision_edge_types: list[EdgeType], + batch_size: int, +): + """Run one loader format on a heterogeneous graph and return its label pairs. + + Results are keyed by ``str(edge_type)`` so they cross the process boundary as + plain data. + """ + create_test_process_group() + loader = DistABLPLoader( + dataset=dataset, + num_neighbors=[2, 2], + input_nodes=input_nodes, + batch_size=batch_size, + pin_memory_device=torch.device("cpu"), + supervision_edge_type=supervision_edge_types, + use_edge_index_output=use_edge_index_output, + ) + positive_pairs: dict[tuple[str, ...], list[tuple[int, int]]] = defaultdict(list) + negative_pairs: dict[tuple[str, ...], list[tuple[int, int]]] = defaultdict(list) + for datum in loader: + assert isinstance(datum, HeteroData) + # Several supervision edge types keep y_positive / y_negative in their + # dict-of-edge-type form rather than collapsing to a bare value. + _accumulate_heterogeneous_pairs( + datum, datum.y_positive, use_edge_index_output, positive_pairs + ) + _accumulate_heterogeneous_pairs( + datum, datum.y_negative, use_edge_index_output, negative_pairs + ) + return_dict[use_edge_index_output] = ( + {edge_type: sorted(pairs) for edge_type, pairs in positive_pairs.items()}, + {edge_type: sorted(pairs) for edge_type, pairs in negative_pairs.items()}, + ) + shutdown_rpc() + + class DistABLPLoaderTest(TestCase): def tearDown(self): if torch.distributed.is_initialized(): @@ -556,6 +727,216 @@ def test_ablp_dataloader( ), ) + @parameterized.expand( + [ + param( + "positive and negative", + labeled_edges={ + _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]), + _NEGATIVE_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 15], [13, 16, 14, 17]] + ), + }, + input_nodes=torch.tensor([10, 15]), + batch_size=2, + has_negatives=True, + ), + param( + "positive only", + labeled_edges={_POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]])}, + input_nodes=torch.tensor([10, 15]), + batch_size=2, + has_negatives=False, + ), + # Anchor 11 has message-passing edges (11 -> {13, 17}) but is the + # source of NO positive-label edge, so its positive-label row is + # all-padding and y_positive[11] is a guaranteed-empty tensor. This + # exercises the empty-anchor branch end-to-end for both outputs. + param( + "guaranteed empty positive anchor", + labeled_edges={ + _POSITIVE_EDGE_TYPE: torch.tensor([[10, 15], [15, 16]]), + _NEGATIVE_EDGE_TYPE: torch.tensor( + [[10, 10, 11, 15], [13, 16, 14, 17]] + ), + }, + input_nodes=torch.tensor([10, 11, 15]), + batch_size=3, + has_negatives=True, + ), + ] + ) + def test_edge_index_output_matches_dict_output( + self, _, labeled_edges, input_nodes, batch_size, has_negatives + ): + """Both output formats contain the same global anchor-label pairs.""" + 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(int(e_idx.max().item() + 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, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir="out") + dataset.build(partition_output=partition_output) + + manager = mp.Manager() + return_dict = manager.dict() + for use_edge_index_output in (False, True): + mp.spawn( + fn=_collect_homogeneous_labels, + args=( + return_dict, + use_edge_index_output, + dataset, + input_nodes, + batch_size, + has_negatives, + ), + ) + self.assertEqual(return_dict[False][0], return_dict[True][0]) + self.assertEqual(return_dict[False][1], return_dict[True][1]) + + @parameterized.expand( + [ + param( + "edge_dir=out", + edge_dir="out", + edge_index={ + _A_TO_B: torch.tensor([[10, 10], [11, 12]]), + message_passing_to_positive_label(_A_TO_B): torch.tensor( + [[10, 10], [13, 14]] + ), + message_passing_to_negative_label(_A_TO_B): torch.tensor( + [[10, 10], [15, 16]] + ), + _A_TO_C: torch.tensor([[10, 10], [20, 21]]), + message_passing_to_positive_label(_A_TO_C): torch.tensor( + [[10, 10], [22, 23]] + ), + message_passing_to_negative_label(_A_TO_C): torch.tensor( + [[10, 10], [24, 25]] + ), + }, + ), + # edge_dir="in" stores the supervision edge types reversed, so their + # dst node type is the anchor type rather than the supervision type, + # while the label edge types reaching collation stay outward-facing. + # A loader that resolves supervision nodes off the former instead of + # the latter fails here and passes the edge_dir="out" case above. + param( + "edge_dir=in", + edge_dir="in", + edge_index={ + _B_TO_A: torch.tensor([[11, 12], [10, 10]]), + message_passing_to_positive_label(_B_TO_A): torch.tensor( + [[13, 14], [10, 10]] + ), + message_passing_to_negative_label(_B_TO_A): torch.tensor( + [[15, 16], [10, 10]] + ), + _C_TO_A: torch.tensor([[20, 21], [10, 10]]), + message_passing_to_positive_label(_C_TO_A): torch.tensor( + [[22, 23], [10, 10]] + ), + message_passing_to_negative_label(_C_TO_A): torch.tensor( + [[24, 25], [10, 10]] + ), + }, + ), + ] + ) + def test_heterogeneous_edge_index_output_matches_dict_output( + self, + _, + edge_dir: Literal["in", "out"], + edge_index: dict[EdgeType, torch.Tensor], + ): + """Both output formats agree on a heterogeneous graph, for either edge_dir. + + Anchors are node type ``a``; supervision nodes are ``b`` and ``c``. Because + the anchor and supervision node maps are distinct, a label remapped against + the wrong node type yields a wrong global id here instead of going unnoticed. + """ + nodes: dict[NodeType, list[torch.Tensor]] = defaultdict(list) + for edge_type, edge_idx in edge_index.items(): + nodes[edge_type[0]].append(edge_idx[0]) + nodes[edge_type[2]].append(edge_idx[1]) + partition_output = PartitionOutput( + node_partition_book={ + node_type: torch.zeros(int(torch.cat(node_ids).max().item() + 1)) + for node_type, node_ids in nodes.items() + }, + edge_partition_book={ + e_type: torch.zeros(int(e_idx.max().item() + 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, + partitioned_node_labels=None, + ) + dataset = DistDataset(rank=0, world_size=1, edge_dir=edge_dir) + dataset.build(partition_output=partition_output) + + manager = mp.Manager() + return_dict = manager.dict() + for use_edge_index_output in (False, True): + mp.spawn( + fn=_collect_heterogeneous_labels, + args=( + return_dict, + use_edge_index_output, + dataset, + (_A, torch.tensor([10])), + [_A_TO_B, _A_TO_C], + 1, # batch_size + ), + ) + # Both formats resolve to the same global ids... + self.assertEqual(dict(return_dict[False][0]), dict(return_dict[True][0])) + self.assertEqual(dict(return_dict[False][1]), dict(return_dict[True][1])) + # ...and did so over labels that actually exist, so equality above cannot + # hold vacuously. + self.assertEqual( + dict(return_dict[True][0]), + { + _edge_type_key(_A_TO_B): [(10, 13), (10, 14)], + _edge_type_key(_A_TO_C): [(10, 22), (10, 23)], + }, + ) + self.assertEqual( + dict(return_dict[True][1]), + { + _edge_type_key(_A_TO_B): [(10, 15), (10, 16)], + _edge_type_key(_A_TO_C): [(10, 24), (10, 25)], + }, + ) + def test_cora_supervised(self): create_test_process_group() cora_supervised_info = get_mocked_dataset_artifact_metadata()[ diff --git a/tests/unit/distributed/label_remap_cuda_device_test.py b/tests/unit/distributed/label_remap_cuda_device_test.py new file mode 100644 index 000000000..53ac1ff07 --- /dev/null +++ b/tests/unit/distributed/label_remap_cuda_device_test.py @@ -0,0 +1,60 @@ +"""CUDA device-placement regression test for ABLP label remapping.""" + +import unittest + +import torch + +from gigl.distributed.utils.ablp import remap_labels_to_local_edge_indices +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import message_passing_to_positive_label +from tests.test_assets.test_case import TestCase + +_USER = NodeType("user") +_STORY = NodeType("story") +_USER_TO_STORY = EdgeType(_USER, Relation("to"), _STORY) + + +def _inputs( + device: torch.device, +) -> tuple[dict[NodeType, torch.Tensor], dict[EdgeType, torch.Tensor]]: + """Build an unsorted node map and padded labels on one device.""" + node_map = {_STORY: torch.tensor([15, 10, 16, 11, 12], device=device)} + positive_labels = { + message_passing_to_positive_label(_USER_TO_STORY): torch.tensor( + [[15, 15], [16, -1], [-1, -1]], + dtype=torch.long, + device=device, + ) + } + return node_map, positive_labels + + +@unittest.skipUnless(torch.cuda.is_available(), "requires a CUDA device") +class LabelRemapCudaDeviceTest(TestCase): + def test_cuda_output_matches_cpu(self) -> None: + cpu_node_map, cpu_positive_labels = _inputs(torch.device("cpu")) + expected_positive, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=cpu_node_map, + positive_labels_by_edge_type=cpu_positive_labels, + negative_labels_by_edge_type={}, + ) + + cuda_node_map, cuda_positive_labels = _inputs(torch.device("cuda")) + actual_positive, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=cuda_node_map, + positive_labels_by_edge_type=cuda_positive_labels, + negative_labels_by_edge_type={}, + ) + + self.assertEqual(set(actual_positive.keys()), set(expected_positive.keys())) + for edge_type, expected_edge_index in expected_positive.items(): + actual_edge_index = actual_positive[edge_type] + self.assertEqual(actual_edge_index.device.type, "cuda") + torch.testing.assert_close( + actual_edge_index.cpu(), + expected_edge_index, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/distributed/utils/ablp_test.py b/tests/unit/distributed/utils/ablp_test.py new file mode 100644 index 000000000..a65726d65 --- /dev/null +++ b/tests/unit/distributed/utils/ablp_test.py @@ -0,0 +1,285 @@ +"""Tests for vectorized ABLP label remapping.""" + +import unittest + +import torch +from parameterized import param, parameterized +from torch_geometric.typing import EdgeType as PyGEdgeType + +from gigl.distributed.utils.ablp import ( + label_edge_index_to_dict, + remap_labels_to_local_edge_indices, +) +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.types.graph import ( + DEFAULT_HOMOGENEOUS_EDGE_TYPE, + DEFAULT_HOMOGENEOUS_NODE_TYPE, + message_passing_to_negative_label, + message_passing_to_positive_label, +) +from tests.test_assets.test_case import TestCase + +_USER = NodeType("user") +_STORY = NodeType("story") +_USER_TO_STORY = EdgeType(_USER, Relation("to"), _STORY) +_A = NodeType("a") +_B = NodeType("b") +_C = NodeType("c") +_A_TO_B = EdgeType(_A, Relation("to"), _B) +_A_TO_C = EdgeType(_A, Relation("to"), _C) + + +def _positive_label_edge_type(edge_type: EdgeType) -> EdgeType: + return message_passing_to_positive_label(edge_type) + + +def _negative_label_edge_type(edge_type: EdgeType) -> EdgeType: + return message_passing_to_negative_label(edge_type) + + +def _assert_label_sets_equal( + actual: dict[PyGEdgeType, torch.Tensor], + expected: dict[PyGEdgeType, dict[int, list[int]]], +) -> None: + assert set(actual.keys()) == set(expected.keys()) + for edge_type, expected_by_anchor in expected.items(): + actual_by_anchor = label_edge_index_to_dict( + label_edge_index=actual[edge_type], + num_anchors=len(expected_by_anchor), + ) + assert set(actual_by_anchor.keys()) == set(expected_by_anchor.keys()) + for anchor, expected_labels in expected_by_anchor.items(): + actual_labels = actual_by_anchor[anchor] + assert actual_labels.dtype == torch.long + assert sorted(actual_labels.tolist()) == sorted(expected_labels) + + +class LabelEdgeIndexToDictTest(TestCase): + def test_preserves_empty_and_multi_label_anchors(self) -> None: + label_edge_index = torch.tensor([[0, 2, 2], [5, 7, 8]]) + + labels_by_anchor = label_edge_index_to_dict( + label_edge_index=label_edge_index, + num_anchors=3, + ) + + self.assertEqual(set(labels_by_anchor.keys()), {0, 1, 2}) + torch.testing.assert_close( + labels_by_anchor[0], torch.tensor([5], dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[1], torch.empty(0, dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[2], torch.tensor([7, 8], dtype=torch.long) + ) + + def test_all_anchors_empty(self) -> None: + labels_by_anchor = label_edge_index_to_dict( + label_edge_index=torch.empty((2, 0), dtype=torch.long), + num_anchors=2, + ) + + self.assertEqual(set(labels_by_anchor.keys()), {0, 1}) + torch.testing.assert_close( + labels_by_anchor[0], torch.empty(0, dtype=torch.long) + ) + torch.testing.assert_close( + labels_by_anchor[1], torch.empty(0, dtype=torch.long) + ) + + +class RemapLabelsToLocalEdgeIndicesTest(TestCase): + @parameterized.expand( + [ + param( + "sorted_present_empty_and_padded", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, -1], [15, 16], [-1, -1], [99, -1]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [5], 1: [5, 6], 2: [], 3: []}}, + expected_negative={}, + ), + param( + "duplicate_labels_preserve_multiplicity", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, 15], [11, 11]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [5, 5], 1: [1, 1]}}, + expected_negative={}, + ), + param( + "default_homogeneous_keying", + local_id_to_global_id_by_node_type={ + DEFAULT_HOMOGENEOUS_NODE_TYPE: torch.tensor([20, 10, 30, 11, 15]) + }, + positive_labels={ + message_passing_to_positive_label( + DEFAULT_HOMOGENEOUS_EDGE_TYPE + ): torch.tensor([[30, 10], [-1, -1]]) + }, + negative_labels={}, + expected_positive={DEFAULT_HOMOGENEOUS_EDGE_TYPE: {0: [2, 1], 1: []}}, + expected_negative={}, + ), + param( + "positive_and_negative_labels", + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15], [16]] + ) + }, + negative_labels={ + _negative_label_edge_type(_USER_TO_STORY): torch.tensor( + [[13, 16], [17, -1]] + ) + }, + expected_positive={_USER_TO_STORY: {0: [5], 1: [6]}}, + expected_negative={_USER_TO_STORY: {0: [3, 6], 1: [7]}}, + ), + param( + "multiple_supervision_edge_types", + local_id_to_global_id_by_node_type={ + _B: torch.tensor([11, 12, 13, 14, 15, 16]), + _C: torch.tensor([20, 21, 22, 23, 24, 25]), + }, + positive_labels={ + _positive_label_edge_type(_A_TO_B): torch.tensor([[13, 14]]), + _positive_label_edge_type(_A_TO_C): torch.tensor([[22, 23]]), + }, + negative_labels={ + _negative_label_edge_type(_A_TO_B): torch.tensor([[15, 16]]), + _negative_label_edge_type(_A_TO_C): torch.tensor([[24, 25]]), + }, + expected_positive={ + _A_TO_B: {0: [2, 3]}, + _A_TO_C: {0: [2, 3]}, + }, + expected_negative={ + _A_TO_B: {0: [4, 5]}, + _A_TO_C: {0: [4, 5]}, + }, + ), + param( + "all_anchors_empty", + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 11, 12])}, + positive_labels={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[-1, -1], [99, 98]] + ) + }, + negative_labels={}, + expected_positive={_USER_TO_STORY: {0: [], 1: []}}, + expected_negative={}, + ), + ] + ) + def test_matches_constructed_labels( + self, + _, + local_id_to_global_id_by_node_type: dict[NodeType, torch.Tensor], + positive_labels: dict[EdgeType, torch.Tensor], + negative_labels: dict[EdgeType, torch.Tensor], + expected_positive: dict[PyGEdgeType, dict[int, list[int]]], + expected_negative: dict[PyGEdgeType, dict[int, list[int]]], + ) -> None: + positive_edge_indices, negative_edge_indices = ( + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type=local_id_to_global_id_by_node_type, + positive_labels_by_edge_type=positive_labels, + negative_labels_by_edge_type=negative_labels, + ) + ) + + _assert_label_sets_equal(positive_edge_indices, expected_positive) + _assert_label_sets_equal(negative_edge_indices, expected_negative) + + def test_pins_exact_edge_index(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15, 16, 17]) + }, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor( + [[15, -1], [15, 16], [-1, -1]] + ) + }, + negative_labels_by_edge_type={}, + ) + + torch.testing.assert_close( + positive_edge_indices[_USER_TO_STORY], + torch.tensor([[0, 1, 1], [5, 5, 6]], dtype=torch.long), + ) + + def test_unsorted_node_map_returns_local_indices(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([15, 10, 16, 11])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[16, 15]]) + }, + negative_labels_by_edge_type={}, + ) + + labels_by_anchor = label_edge_index_to_dict( + positive_edge_indices[_USER_TO_STORY], num_anchors=1 + ) + self.assertEqual(sorted(labels_by_anchor[0].tolist()), [0, 2]) + + def test_zero_anchor_tensor_yields_no_edge_type_key(self) -> None: + positive_edge_indices, negative_edge_indices = ( + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 11, 12])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.empty((0, 0)) + }, + negative_labels_by_edge_type={}, + ) + ) + + self.assertEqual(positive_edge_indices, {}) + self.assertEqual(negative_edge_indices, {}) + + def test_output_follows_input_device_and_dtype(self) -> None: + positive_edge_indices, _ = remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={ + _STORY: torch.tensor([10, 11, 12, 13, 14, 15]) + }, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[15], [11]]) + }, + negative_labels_by_edge_type={}, + ) + + label_edge_index = positive_edge_indices[_USER_TO_STORY] + self.assertEqual(label_edge_index.device.type, "cpu") + self.assertEqual(label_edge_index.dtype, torch.long) + + def test_duplicate_node_map_raises(self) -> None: + with self.assertRaises(ValueError): + remap_labels_to_local_edge_indices( + local_id_to_global_id_by_node_type={_STORY: torch.tensor([10, 10, 11])}, + positive_labels_by_edge_type={ + _positive_label_edge_type(_USER_TO_STORY): torch.tensor([[10, 11]]) + }, + negative_labels_by_edge_type={}, + ) + + +if __name__ == "__main__": + unittest.main()