diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index b67611b2f..a2ecb6d65 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -22,12 +22,15 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.transforms.graph_transformer import ( + PPR_RELATION_FEATURES_NAME, PPR_WEIGHT_FEATURE_NAME, SequenceAuxiliaryData, TokenInputData, heterodata_to_graph_transformer_input, ) +_L2_NORMALIZE_EPS = 1e-6 + def _get_node_type_positional_encodings( data: torch_geometric.data.hetero_data.HeteroData, @@ -79,6 +82,26 @@ def _build_sinusoidal_sequence_position_table( return position_table +def _zero_initialize_lazy_linear_if_needed(module: nn.Module, input: Tensor) -> None: + """Materialize an uninitialized LazyLinear and make it initially no-op.""" + has_uninitialized_params = getattr(module, "has_uninitialized_params", None) + if not callable(has_uninitialized_params) or not has_uninitialized_params(): + return + + initialize_parameters = getattr(module, "initialize_parameters", None) + if not callable(initialize_parameters): + return + + initialize_parameters(input) + with torch.no_grad(): + weight = getattr(module, "weight", None) + if isinstance(weight, nn.Parameter): + weight.zero_() + bias = getattr(module, "bias", None) + if isinstance(bias, nn.Parameter): + bias.zero_() + + # Supported activation functions for FeedForwardNetwork _ACTIVATION_FNS = { "gelu": nn.GELU, @@ -682,8 +705,11 @@ class GraphTransformerEncoder(nn.Module): anchor_based_input_attr_names: List of anchor-relative attribute names used as token-aligned input features. Sparse graph-level attributes are looked up from ``data`` and ``"ppr_weight"`` resolves to PPR - edge weights in PPR mode. These are projected to ``hid_dim`` and - added to the sequence tokens after sequence construction. + edge weights in PPR mode. The reserved name + ``"ppr_relation_features"`` resolves to additional PPR edge-attr + columns after the first weight column. These are projected to + ``hid_dim`` and added to the sequence tokens after sequence + construction. Example: ``['hop_distance', 'ppr_weight']`` for continuous features, or ``['hop_distance']`` when ``hop_distance`` will be embedded via ``anchor_based_input_embedding_dict``. @@ -858,19 +884,29 @@ def __init__( anchor_bias_attr_names = anchor_based_attention_bias_attr_names or [] anchor_input_attr_names = anchor_based_input_attr_names or [] pairwise_bias_attr_names = pairwise_attention_bias_attr_names or [] - if PPR_WEIGHT_FEATURE_NAME in pairwise_bias_attr_names: + ppr_reserved_feature_names = { + PPR_WEIGHT_FEATURE_NAME, + PPR_RELATION_FEATURES_NAME, + } + if ppr_reserved_feature_names & set(pairwise_bias_attr_names): raise ValueError( - f"'{PPR_WEIGHT_FEATURE_NAME}' is an anchor-relative feature and " + f"{sorted(ppr_reserved_feature_names)} are anchor-relative features and " "cannot be used as pairwise attention bias." ) if ( - PPR_WEIGHT_FEATURE_NAME in anchor_bias_attr_names + anchor_input_attr_names + ppr_reserved_feature_names + & set(anchor_bias_attr_names + anchor_input_attr_names) and sequence_construction_method != "ppr" ): raise ValueError( - "The reserved anchor-relative feature 'ppr_weight' requires " + "Reserved PPR anchor-relative features require " "sequence_construction_method='ppr'." ) + if PPR_RELATION_FEATURES_NAME in anchor_bias_attr_names: + raise ValueError( + f"'{PPR_RELATION_FEATURES_NAME}' is a multi-column token-input " + "feature and cannot be used as attention bias." + ) self._sequence_construction_method = sequence_construction_method self._sampling_direction = sampling_direction self._sequence_positional_encoding_type = sequence_positional_encoding_type @@ -931,7 +967,6 @@ def __init__( None, persistent=False, ) - # Per-node-type input projection to hid_dim (like HGT's lin_dict) self._node_projection_dict = nn.ModuleDict( { @@ -1096,14 +1131,14 @@ def forward( if hasattr(data[edge_type], "edge_attr"): projected_data[edge_type].edge_attr = data[edge_type].edge_attr # Copy relative-encoding attributes (e.g., hop_distance stored as sparse matrix) - relative_pe_attr_names = { - attr_name - for attr_name in (self._anchor_based_attention_bias_attr_names or []) - if attr_name != PPR_WEIGHT_FEATURE_NAME - } - relative_pe_attr_names.update(self._anchor_based_input_attr_names or []) - relative_pe_attr_names.update(self._pairwise_attention_bias_attr_names or []) - relative_pe_attr_names.discard(PPR_WEIGHT_FEATURE_NAME) + relative_pe_attr_names = ( + set(self._anchor_based_attention_bias_attr_names or []) + | set(self._anchor_based_input_attr_names or []) + | set(self._pairwise_attention_bias_attr_names or []) + ) + relative_pe_attr_names.difference_update( + {PPR_WEIGHT_FEATURE_NAME, PPR_RELATION_FEATURES_NAME} + ) if relative_pe_attr_names: for attr_name in sorted(relative_pe_attr_names): if hasattr(data, attr_name): @@ -1185,7 +1220,8 @@ def forward( embeddings = self._output_projection(embeddings) if self._should_l2_normalize_embedding_layer_output: - embeddings = F.normalize(embeddings, p=2, dim=-1) + # Default eps=1e-12 underflows to zero in fp16 autocast. + embeddings = F.normalize(embeddings, p=2, dim=-1, eps=_L2_NORMALIZE_EPS) return embeddings @@ -1268,11 +1304,21 @@ def _build_token_input_contribution( "sequence auxiliary data." ) continuous_feature_parts.append(token_input_features[attr_name]) + continuous_features = torch.cat(continuous_feature_parts, dim=-1).to( + sequences.dtype + ) + continuous_features = torch.nan_to_num( + continuous_features, + nan=0.0, + posinf=1.0, + neginf=0.0, + ) + _zero_initialize_lazy_linear_if_needed( + module=self._token_input_projection, + input=continuous_features, + ) token_contribution = token_contribution + ( - self._token_input_projection( - torch.cat(continuous_feature_parts, dim=-1).to(sequences.dtype) - ) - * valid_token_mask + self._token_input_projection(continuous_features) * valid_token_mask ) return token_contribution @@ -1410,6 +1456,41 @@ def _build_attention_bias( return attn_bias + def _readout_from_encoded_sequences( + self, + x: Tensor, + valid_mask: Tensor, + ) -> Tensor: + anchor = x[:, 0, :].unsqueeze(1) + if self._readout_mode == "anchor_only": + return anchor.squeeze(1) + + # Historical readout: anchor token plus attention-weighted neighbors. + neighbors = x[:, 1:, :] + neighbor_valid_mask = valid_mask[:, 1:] + seq_minus_one = neighbors.size(1) + + if seq_minus_one == 0: + return anchor.squeeze(1) + + anchor_expanded = anchor.expand(-1, seq_minus_one, -1) + if self._readout_attention is None: + raise ValueError("Readout attention layer is not initialized.") + readout_scores = self._readout_attention( + torch.cat([anchor_expanded, neighbors], dim=-1) + ) + readout_scores = readout_scores.masked_fill( + ~neighbor_valid_mask.unsqueeze(-1), + torch.finfo(readout_scores.dtype).min, + ) + readout_weights = F.softmax(readout_scores, dim=1) + readout_weights = torch.nan_to_num(readout_weights, nan=0.0) + readout_weights = readout_weights * neighbor_valid_mask.unsqueeze(-1).to( + readout_weights.dtype + ) + neighbor_aggregation = (neighbors * readout_weights).sum(dim=1, keepdim=True) + return (anchor + neighbor_aggregation).squeeze(1) + def _encode_and_readout( self, sequences: Tensor, @@ -1443,41 +1524,4 @@ def _encode_and_readout( x = self._final_norm(x) x = x * valid_mask.unsqueeze(-1).to(x.dtype) - # Readout: anchor (position 0) + attention-weighted neighbor aggregation - anchor = x[:, 0, :].unsqueeze(1) # (batch, 1, hid_dim) - if self._readout_mode == "anchor_only": - return anchor.squeeze(1) - - neighbors = x[:, 1:, :] # (batch, seq-1, hid_dim) - neighbor_valid_mask = valid_mask[:, 1:] - seq_minus_one = neighbors.size(1) - - if seq_minus_one == 0: - return anchor.squeeze(1) - - # Expand anchor to match neighbor dimension for concatenation - anchor_expanded = anchor.expand(-1, seq_minus_one, -1) - - # Compute attention scores over neighbors - if self._readout_attention is None: - raise ValueError("Readout attention layer is not initialized.") - readout_scores = self._readout_attention( - torch.cat([anchor_expanded, neighbors], dim=-1) - ) # (batch, seq-1, 1) - readout_scores = readout_scores.masked_fill( - ~neighbor_valid_mask.unsqueeze(-1), - torch.finfo(readout_scores.dtype).min, - ) - readout_weights = F.softmax(readout_scores, dim=1) # (batch, seq-1, 1) - readout_weights = torch.nan_to_num(readout_weights, nan=0.0) - readout_weights = readout_weights * neighbor_valid_mask.unsqueeze(-1).to( - readout_weights.dtype - ) - - neighbor_aggregation = (neighbors * readout_weights).sum( - dim=1, keepdim=True - ) # (batch, 1, hid_dim) - - output = (anchor + neighbor_aggregation).squeeze(1) # (batch, hid_dim) - - return output + return self._readout_from_encoded_sequences(x=x, valid_mask=valid_mask) diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index e17ed87ff..d22a29f2f 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -79,6 +79,7 @@ class SequenceAuxiliaryData(TypedDict): PPR_WEIGHT_FEATURE_NAME = "ppr_weight" +PPR_RELATION_FEATURES_NAME = "ppr_relation_features" class _TokenOccurrenceIndex(NamedTuple): @@ -131,7 +132,9 @@ def heterodata_to_graph_transformer_input( sequence_construction_method: Strategy used to build per-anchor sequences. ``"khop"`` performs the existing k-hop expansion over the sampled graph. ``"ppr"`` uses outgoing ``(anchor_type, "ppr", neighbor_type)`` edges, - sorted by descending PPR weight from ``edge_attr``. (default: ``"khop"``) + sorted by descending PPR weight for scalar-PPR edge attrs. Typed-PPR + multi-column edge attrs preserve sampler order, which may encode + relation blocks. (default: ``"khop"``) include_anchor_first: If True, anchor node is always first in sequence. padding_value: Value to use for padding (default: 0.0). anchor_based_attention_bias_attr_names: List of anchor-relative feature @@ -142,8 +145,10 @@ def heterodata_to_graph_transformer_input( anchor_based_input_attr_names: List of anchor-relative attribute names returned as token-aligned model-input features. Sparse graph-level attributes are looked up from ``data`` and ``"ppr_weight"`` resolves - to PPR edge weights in PPR sequence mode. - Example: ['hop_distance', 'ppr_weight']. + to PPR edge weights in PPR sequence mode. The reserved name + ``"ppr_relation_features"`` resolves to any additional PPR edge-attr + columns after the first weight column. + Example: ['hop_distance', 'ppr_weight', 'ppr_relation_features']. pairwise_attention_bias_attr_names: List of pairwise feature names used as attention bias. These must correspond to sparse graph-level attributes on ``data``. Example: ['pairwise_distance']. @@ -176,7 +181,7 @@ def heterodata_to_graph_transformer_input( storing ``(batch_idx, row_pos, col_pos)`` coordinates for nonmissing pairwise entries ``"token_input"`` as a dict mapping attribute name to a - ``(batch, seq, 1)`` tensor, or None + ``(batch, seq, attr_dim)`` tensor, or None Raises: ValueError: If node types have different feature dimensions. @@ -216,9 +221,15 @@ def heterodata_to_graph_transformer_input( anchor_input_attr_names = anchor_based_input_attr_names or [] pairwise_bias_attr_names = pairwise_attention_bias_attr_names or [] - if PPR_WEIGHT_FEATURE_NAME in pairwise_bias_attr_names: + ppr_reserved_feature_names = { + PPR_WEIGHT_FEATURE_NAME, + PPR_RELATION_FEATURES_NAME, + } + requested_ppr_feature_names = set(anchor_bias_attr_names + anchor_input_attr_names) + + if ppr_reserved_feature_names & set(pairwise_bias_attr_names): raise ValueError( - f"'{PPR_WEIGHT_FEATURE_NAME}' is an anchor-relative feature and cannot " + f"{sorted(ppr_reserved_feature_names)} are anchor-relative features and cannot " "be used as pairwise attention bias." ) @@ -234,14 +245,20 @@ def heterodata_to_graph_transformer_input( ) if ( - PPR_WEIGHT_FEATURE_NAME in anchor_bias_attr_names + anchor_input_attr_names + ppr_reserved_feature_names & requested_ppr_feature_names and sequence_construction_method != "ppr" ): raise ValueError( - "The reserved anchor-relative feature 'ppr_weight' requires " + "Reserved PPR anchor-relative features require " "sequence_construction_method='ppr'." ) + if PPR_RELATION_FEATURES_NAME in anchor_bias_attr_names: + raise ValueError( + f"'{PPR_RELATION_FEATURES_NAME}' is a multi-column token-input feature " + "and cannot be used as attention bias." + ) + if sequence_construction_method == "ppr": _validate_ppr_sequence_input(data) @@ -274,6 +291,7 @@ def heterodata_to_graph_transformer_input( anchor_indices = offset + anchor_local_indices ppr_weight_sequences: Optional[Tensor] = None + ppr_relation_feature_sequences: Optional[Tensor] = None if sequence_construction_method == "khop": homo_edge_index = homo_data.edge_index # (2, num_edges) if sampling_direction == "in": @@ -299,6 +317,7 @@ def heterodata_to_graph_transformer_input( node_index_sequences, valid_mask, ppr_weight_sequences, + ppr_relation_feature_sequences, ) = _build_sequence_layout_from_ppr_edges( homo_data=homo_data, anchor_indices=anchor_indices, @@ -310,6 +329,9 @@ def heterodata_to_graph_transformer_input( PPR_WEIGHT_FEATURE_NAME in anchor_bias_attr_names + anchor_input_attr_names ), + return_relation_features=( + PPR_RELATION_FEATURES_NAME in anchor_input_attr_names + ), ) else: raise ValueError( @@ -321,7 +343,7 @@ def heterodata_to_graph_transformer_input( { attr_name for attr_name in (anchor_bias_attr_names + anchor_input_attr_names) - if attr_name != PPR_WEIGHT_FEATURE_NAME + if attr_name not in ppr_reserved_feature_names } ) anchor_based_matrices = _get_sparse_feature_matrices( @@ -375,12 +397,14 @@ def heterodata_to_graph_transformer_input( available_anchor_attr_names=anchor_matrix_attr_names, requested_anchor_attr_names=anchor_bias_attr_names, ppr_weight_sequences=ppr_weight_sequences, + ppr_relation_feature_sequences=ppr_relation_feature_sequences, ) token_input_features = _compose_anchor_feature_dict( anchor_relative_feature_sequences=anchor_relative_feature_sequences, available_anchor_attr_names=anchor_matrix_attr_names, requested_anchor_attr_names=anchor_input_attr_names, ppr_weight_sequences=ppr_weight_sequences, + ppr_relation_feature_sequences=ppr_relation_feature_sequences, ) return ( @@ -450,6 +474,7 @@ def _compose_anchor_feature_tensor( available_anchor_attr_names: list[str], requested_anchor_attr_names: list[str], ppr_weight_sequences: Optional[Tensor], + ppr_relation_feature_sequences: Optional[Tensor], ) -> Optional[Tensor]: if not requested_anchor_attr_names: return None @@ -460,25 +485,14 @@ def _compose_anchor_feature_tensor( } for attr_name in requested_anchor_attr_names: - if attr_name == PPR_WEIGHT_FEATURE_NAME: - if ppr_weight_sequences is None: - raise ValueError( - f"Requested '{PPR_WEIGHT_FEATURE_NAME}' but it was not computed." - ) - feature_parts.append(ppr_weight_sequences) - continue - - if anchor_relative_feature_sequences is None: - raise ValueError( - "Anchor-relative features were requested but not computed." - ) - if attr_name not in feature_index_by_name: - raise ValueError( - f"Anchor-relative feature '{attr_name}' was requested but not found." - ) - feature_idx = feature_index_by_name[attr_name] feature_parts.append( - anchor_relative_feature_sequences[..., feature_idx : feature_idx + 1] + _resolve_anchor_feature_sequence( + attr_name=attr_name, + anchor_relative_feature_sequences=anchor_relative_feature_sequences, + feature_index_by_name=feature_index_by_name, + ppr_weight_sequences=ppr_weight_sequences, + ppr_relation_feature_sequences=ppr_relation_feature_sequences, + ) ) return torch.cat(feature_parts, dim=-1) @@ -489,6 +503,7 @@ def _compose_anchor_feature_dict( available_anchor_attr_names: list[str], requested_anchor_attr_names: list[str], ppr_weight_sequences: Optional[Tensor], + ppr_relation_feature_sequences: Optional[Tensor], ) -> Optional[TokenInputData]: if not requested_anchor_attr_names: return None @@ -499,28 +514,45 @@ def _compose_anchor_feature_dict( } for attr_name in requested_anchor_attr_names: - if attr_name == PPR_WEIGHT_FEATURE_NAME: - if ppr_weight_sequences is None: - raise ValueError( - f"Requested '{PPR_WEIGHT_FEATURE_NAME}' but it was not computed." - ) - feature_dict[attr_name] = ppr_weight_sequences - continue + feature_dict[attr_name] = _resolve_anchor_feature_sequence( + attr_name=attr_name, + anchor_relative_feature_sequences=anchor_relative_feature_sequences, + feature_index_by_name=feature_index_by_name, + ppr_weight_sequences=ppr_weight_sequences, + ppr_relation_feature_sequences=ppr_relation_feature_sequences, + ) - if anchor_relative_feature_sequences is None: + return feature_dict + + +def _resolve_anchor_feature_sequence( + attr_name: str, + anchor_relative_feature_sequences: Optional[Tensor], + feature_index_by_name: dict[str, int], + ppr_weight_sequences: Optional[Tensor], + ppr_relation_feature_sequences: Optional[Tensor], +) -> Tensor: + if attr_name == PPR_WEIGHT_FEATURE_NAME: + if ppr_weight_sequences is None: raise ValueError( - "Anchor-relative features were requested but not computed." + f"Requested '{PPR_WEIGHT_FEATURE_NAME}' but it was not computed." ) - if attr_name not in feature_index_by_name: + return ppr_weight_sequences + if attr_name == PPR_RELATION_FEATURES_NAME: + if ppr_relation_feature_sequences is None: raise ValueError( - f"Anchor-relative feature '{attr_name}' was requested but not found." + f"Requested '{PPR_RELATION_FEATURES_NAME}' but it was not computed." ) - feature_idx = feature_index_by_name[attr_name] - feature_dict[attr_name] = anchor_relative_feature_sequences[ - ..., feature_idx : feature_idx + 1 - ] + return ppr_relation_feature_sequences - return feature_dict + if anchor_relative_feature_sequences is None: + raise ValueError("Anchor-relative features were requested but not computed.") + if attr_name not in feature_index_by_name: + raise ValueError( + f"Anchor-relative feature '{attr_name}' was requested but not found." + ) + feature_idx = feature_index_by_name[attr_name] + return anchor_relative_feature_sequences[..., feature_idx : feature_idx + 1] def _build_sequence_layout_from_sparse_neighbors( @@ -627,13 +659,15 @@ def _build_sequence_layout_from_ppr_edges( num_nodes: int, device: torch.device, return_edge_weights: bool = False, -) -> tuple[Tensor, Tensor, Optional[Tensor]]: + return_relation_features: bool = False, +) -> tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]: """Build sequences directly from outgoing PPR edges for each anchor. The sequence order is: 1. Anchor node first, when ``include_anchor_first`` is True. - 2. Destination nodes reachable by outgoing ``"ppr"`` edges from that anchor, - sorted by descending PPR weight. + 2. Destination nodes reachable by outgoing ``"ppr"`` edges from that anchor. + Scalar-PPR edge attrs are sorted by descending PPR weight; typed-PPR + multi-column edge attrs preserve sampler order. """ batch_size = anchor_indices.size(0) node_index_sequences = torch.full( @@ -654,6 +688,7 @@ def _build_sequence_layout_from_ppr_edges( dtype=torch.float, device=device, ) + ppr_relation_feature_sequences = None if include_anchor_first and max_seq_len > 0: node_index_sequences[:, 0] = anchor_indices @@ -663,26 +698,45 @@ def _build_sequence_layout_from_ppr_edges( start_pos = 0 if start_pos >= max_seq_len: - return node_index_sequences, valid_mask, ppr_weight_sequences + return ( + node_index_sequences, + valid_mask, + ppr_weight_sequences, + ppr_relation_feature_sequences, + ) if not hasattr(homo_data, "edge_attr") or homo_data.edge_attr is None: raise ValueError( "sequence_construction_method='ppr' requires homogeneous edge_attr weights." ) - edge_weights = homo_data.edge_attr - if edge_weights.dim() == 2: - if edge_weights.size(1) != 1: + edge_features = torch.nan_to_num( + homo_data.edge_attr.float(), + nan=0.0, + posinf=1.0, + neginf=0.0, + ).clamp_(min=0.0, max=1.0) + if edge_features.dim() == 1: + edge_features = edge_features.unsqueeze(1) + elif edge_features.dim() != 2: + raise ValueError( + f"PPR edge features must be 1D or 2D, got {tuple(edge_features.shape)}." + ) + if edge_features.size(1) < 1: + raise ValueError("PPR edge features must contain at least one weight column.") + if return_relation_features: + if edge_features.size(1) == 1: raise ValueError( - "PPR edge weights must be 1D or shape [N, 1], " - f"got {tuple(edge_weights.shape)}." + f"Requested '{PPR_RELATION_FEATURES_NAME}' but PPR edge_attr only " + "contains the scalar weight column." ) - edge_weights = edge_weights.squeeze(1) - elif edge_weights.dim() != 1: - raise ValueError( - "PPR edge weights must be 1D or shape [N, 1], " - f"got {tuple(edge_weights.shape)}." + ppr_relation_feature_sequences = torch.zeros( + (batch_size, max_seq_len, edge_features.size(1) - 1), + dtype=torch.float, + device=device, ) + edge_weights = edge_features[:, 0] + relation_features = edge_features[:, 1:] if edge_features.size(1) > 1 else None anchor_batch_index_by_homo_idx = torch.full( (num_nodes,), @@ -699,32 +753,54 @@ def _build_sequence_layout_from_ppr_edges( anchor_batch_idx = anchor_batch_index_by_homo_idx[src_idx] keep = anchor_batch_idx >= 0 if not keep.any(): - return node_index_sequences, valid_mask, ppr_weight_sequences + return ( + node_index_sequences, + valid_mask, + ppr_weight_sequences, + ppr_relation_feature_sequences, + ) all_anchor_batch_idx = anchor_batch_idx[keep] all_dst_idx = dst_idx[keep] all_weights = edge_weights[keep] + all_relation_features = ( + relation_features[keep] if relation_features is not None else None + ) if include_anchor_first: keep = all_dst_idx != anchor_indices[all_anchor_batch_idx] if not keep.any(): - return node_index_sequences, valid_mask, ppr_weight_sequences + return ( + node_index_sequences, + valid_mask, + ppr_weight_sequences, + ppr_relation_feature_sequences, + ) all_anchor_batch_idx = all_anchor_batch_idx[keep] all_dst_idx = all_dst_idx[keep] all_weights = all_weights[keep] + all_relation_features = ( + all_relation_features[keep] if all_relation_features is not None else None + ) - # Flattened COO edges can be laid out in one pass by sorting first on weight - # and then stably on anchor batch id, which preserves descending-weight order - # within each anchor group without a Python loop. - weight_order = torch.argsort(all_weights, descending=True, stable=True) - all_anchor_batch_idx = all_anchor_batch_idx[weight_order] - all_dst_idx = all_dst_idx[weight_order] - all_weights = all_weights[weight_order] + # Regular scalar-PPR edges are sorted by weight here. Typed-PPR emits + # multi-column edge_attr and uses sampler order to preserve relation blocks, + # so keep that order within each anchor group. + if all_relation_features is None: + weight_order = torch.argsort(all_weights, descending=True, stable=True) + all_anchor_batch_idx = all_anchor_batch_idx[weight_order] + all_dst_idx = all_dst_idx[weight_order] + all_weights = all_weights[weight_order] batch_order = torch.argsort(all_anchor_batch_idx, stable=True) sorted_batch_idx = all_anchor_batch_idx[batch_order] sorted_dst_idx = all_dst_idx[batch_order] sorted_weights = all_weights[batch_order] + sorted_relation_features = ( + all_relation_features[batch_order] + if all_relation_features is not None + else None + ) n = sorted_batch_idx.size(0) is_group_start = torch.zeros(n, dtype=torch.long, device=device) @@ -741,6 +817,11 @@ def _build_sequence_layout_from_ppr_edges( valid_positions = positions[valid] valid_dst_idx = sorted_dst_idx[valid] valid_weights = sorted_weights[valid] + valid_relation_features = ( + sorted_relation_features[valid] + if sorted_relation_features is not None + else None + ) node_index_sequences[valid_batch_idx, valid_positions] = valid_dst_idx valid_mask[valid_batch_idx, valid_positions] = True @@ -748,8 +829,20 @@ def _build_sequence_layout_from_ppr_edges( ppr_weight_sequences[valid_batch_idx, valid_positions, 0] = ( valid_weights.float() ) + if ( + ppr_relation_feature_sequences is not None + and valid_relation_features is not None + ): + ppr_relation_feature_sequences[valid_batch_idx, valid_positions] = ( + valid_relation_features.float() + ) - return node_index_sequences, valid_mask, ppr_weight_sequences + return ( + node_index_sequences, + valid_mask, + ppr_weight_sequences, + ppr_relation_feature_sequences, + ) def _gather_sequences_from_node_indices( diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index 70cdf7eca..d9e451de6 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -14,6 +14,7 @@ GraphTransformerEncoderLayer, ) from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation +from gigl.transforms.graph_transformer import PPR_RELATION_FEATURES_NAME from tests.test_assets.test_case import TestCase @@ -1285,6 +1286,34 @@ def test_forward_supports_mixed_embedded_and_continuous_token_input_features( torch.allclose(base_embeddings, augmented_embeddings, atol=1e-6) ) + def test_forward_supports_ppr_relation_token_input_features(self) -> None: + data = _create_user_graph_with_ppr_edges() + data["user", "ppr", "user"].edge_attr = torch.tensor( + [ + [0.9, 1.0, 0.0], + [0.4, 0.0, 1.0], + [0.7, 1.0, 1.0], + ] + ) + ppr_edge_type = EdgeType(self._node_type, Relation("ppr"), self._node_type) + + encoder = self._create_encoder( + edge_type_to_feat_dim_map={ppr_edge_type: 0}, + sequence_construction_method="ppr", + anchor_based_input_attr_names=[PPR_RELATION_FEATURES_NAME], + ) + encoder.eval() + + with torch.no_grad(): + embeddings = encoder( + data=data, + anchor_node_type=self._node_type, + device=self._device, + ) + + self.assertEqual(embeddings.shape, (3, 6)) + self.assertFalse(torch.isnan(embeddings).any()) + class TestFeedForwardNetwork(TestCase): """Tests for FeedForwardNetwork with various activations.""" diff --git a/tests/unit/transforms/graph_transformer_test.py b/tests/unit/transforms/graph_transformer_test.py index d1014045f..a652a648e 100644 --- a/tests/unit/transforms/graph_transformer_test.py +++ b/tests/unit/transforms/graph_transformer_test.py @@ -11,6 +11,8 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.transforms.graph_transformer import ( + PPR_RELATION_FEATURES_NAME, + PPR_WEIGHT_FEATURE_NAME, _get_k_hop_neighbors_sparse, heterodata_to_graph_transformer_input, ) @@ -636,6 +638,105 @@ def test_ppr_sequence_can_return_token_input_and_attention_bias_features(self): torch.equal(valid_mask[1], torch.tensor([True, True, True, False])) ) + def test_ppr_sequence_can_return_relation_token_input_features(self): + data = create_ppr_sequence_hetero_data() + data["user", "ppr", "item"].edge_attr = torch.tensor( + [ + [0.9, 9.0, 90.0], + [0.6, 6.0, 60.0], + [0.8, 8.0, 80.0], + ] + ) + data["user", "ppr", "user"].edge_attr = torch.tensor( + [ + [0.4, 4.0, 40.0], + [0.3, 3.0, 30.0], + ] + ) + + _, _, sequence_auxiliary_data = heterodata_to_graph_transformer_input( + data=data, + batch_size=2, + max_seq_len=4, + anchor_node_type="user", + sequence_construction_method="ppr", + anchor_based_input_attr_names=[ + PPR_WEIGHT_FEATURE_NAME, + PPR_RELATION_FEATURES_NAME, + ], + ) + + token_input = sequence_auxiliary_data["token_input"] + assert token_input is not None + + self.assertEqual( + set(token_input.keys()), + {PPR_WEIGHT_FEATURE_NAME, PPR_RELATION_FEATURES_NAME}, + ) + self.assertTrue( + torch.allclose( + token_input[PPR_WEIGHT_FEATURE_NAME], + torch.tensor( + [ + [[0.0], [0.9], [0.6], [0.4]], + [[0.0], [0.8], [0.3], [0.0]], + ] + ), + ) + ) + self.assertTrue( + torch.allclose( + token_input[PPR_RELATION_FEATURES_NAME], + torch.tensor( + [ + [[0.0, 0.0], [1.0, 1.0], [1.0, 1.0], [1.0, 1.0]], + [[0.0, 0.0], [1.0, 1.0], [1.0, 1.0], [0.0, 0.0]], + ] + ), + ) + ) + + def test_ppr_sequence_preserves_typed_ppr_sampler_order(self): + data = HeteroData() + data["user"].x = torch.tensor([[0.0], [1.0], [2.0]]) + data["user", "ppr", "user"].edge_index = torch.tensor( + [ + [0, 0], + [1, 2], + ] + ) + data["user", "ppr", "user"].edge_attr = torch.tensor( + [ + [0.2, 1.0, 0.0], + [0.9, 0.0, 1.0], + ] + ) + + sequences, _, sequence_auxiliary_data = heterodata_to_graph_transformer_input( + data=data, + batch_size=1, + max_seq_len=3, + anchor_node_type="user", + sequence_construction_method="ppr", + anchor_based_attention_bias_attr_names=[PPR_WEIGHT_FEATURE_NAME], + ) + + anchor_bias = sequence_auxiliary_data["anchor_bias"] + assert anchor_bias is not None + + self.assertTrue( + torch.allclose( + sequences.squeeze(-1), + torch.tensor([[0.0, 1.0, 2.0]]), + ) + ) + self.assertTrue( + torch.allclose( + anchor_bias.squeeze(-1), + torch.tensor([[0.0, 0.2, 0.9]]), + ) + ) + class TestPyTorchTransformerIntegration(TestCase): """Tests for integration with PyTorch TransformerEncoderLayer."""