From 241f7a5f06a3e5342658d44fd3b860526b31875a Mon Sep 17 00:00:00 2001 From: mkolodner Date: Fri, 10 Jul 2026 19:31:21 +0000 Subject: [PATCH 1/6] Add PPR relation token features to graph transformer --- gigl/nn/graph_transformer.py | 160 ++++++++----- gigl/transforms/graph_transformer.py | 220 +++++++++++++----- tests/unit/nn/graph_transformer_test.py | 29 +++ .../unit/transforms/graph_transformer_test.py | 110 +++++++++ 4 files changed, 398 insertions(+), 121 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index a43c3f5e1..ddfd1b886 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, @@ -1100,8 +1123,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``. @@ -1308,19 +1334,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 @@ -1385,7 +1421,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( { @@ -1551,14 +1586,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): @@ -1641,7 +1676,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 @@ -1724,11 +1760,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 @@ -1874,6 +1920,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, @@ -1938,41 +2019,4 @@ def _encode_and_readout( x = self._final_norm(x) x = x * output_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..faa0f6d26 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 from edge-attr column 0. + Multi-column edge attrs can also expose remaining columns as + ``"ppr_relation_features"``. (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. + PPR edge attrs are sorted by descending weight from column 0. Any + additional edge-attr columns are carried with the same permutation. """ 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,55 @@ 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] + all_relation_features = ( + all_relation_features[weight_order] + if all_relation_features is not None + else None + ) 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 +818,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 +830,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 3684ae4e4..e24276eb2 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -15,6 +15,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 @@ -1959,6 +1960,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..b70bb6a38 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,114 @@ 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_sorts_multi_column_edge_attr_by_weight_column(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_based_input_attr_names=[PPR_RELATION_FEATURES_NAME], + ) + + anchor_bias = sequence_auxiliary_data["anchor_bias"] + token_input = sequence_auxiliary_data["token_input"] + assert anchor_bias is not None + assert token_input is not None + + self.assertTrue( + torch.allclose( + sequences.squeeze(-1), + torch.tensor([[0.0, 2.0, 1.0]]), + ) + ) + self.assertTrue( + torch.allclose( + anchor_bias.squeeze(-1), + torch.tensor([[0.0, 0.9, 0.2]]), + ) + ) + self.assertTrue( + torch.allclose( + token_input[PPR_RELATION_FEATURES_NAME], + torch.tensor([[[0.0, 0.0], [0.0, 1.0], [1.0, 0.0]]]), + ) + ) + class TestPyTorchTransformerIntegration(TestCase): """Tests for integration with PyTorch TransformerEncoderLayer.""" From aa84c675c6e1686f2b5229281bdb0ad0bd406a21 Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 29 Jul 2026 10:37:38 +0000 Subject: [PATCH 2/6] Simplify PPR relation feature handling --- gigl/nn/graph_transformer.py | 30 ------------------- gigl/transforms/graph_transformer.py | 7 +---- .../unit/transforms/graph_transformer_test.py | 4 +-- 3 files changed, 3 insertions(+), 38 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index ddfd1b886..1f78f8826 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -82,26 +82,6 @@ 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, @@ -1763,16 +1743,6 @@ def _build_token_input_contribution( 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(continuous_features) * valid_token_mask ) diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index faa0f6d26..2e520a945 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -710,12 +710,7 @@ def _build_sequence_layout_from_ppr_edges( "sequence_construction_method='ppr' requires homogeneous edge_attr weights." ) - 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) + edge_features = homo_data.edge_attr.float() if edge_features.dim() == 1: edge_features = edge_features.unsqueeze(1) elif edge_features.dim() != 2: diff --git a/tests/unit/transforms/graph_transformer_test.py b/tests/unit/transforms/graph_transformer_test.py index b70bb6a38..ec95b278d 100644 --- a/tests/unit/transforms/graph_transformer_test.py +++ b/tests/unit/transforms/graph_transformer_test.py @@ -689,8 +689,8 @@ def test_ppr_sequence_can_return_relation_token_input_features(self): 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]], + [[0.0, 0.0], [9.0, 90.0], [6.0, 60.0], [4.0, 40.0]], + [[0.0, 0.0], [8.0, 80.0], [3.0, 30.0], [0.0, 0.0]], ] ), ) From 3ba954e646124f7294621686f98ed9d63b332acf Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 29 Jul 2026 18:55:50 +0000 Subject: [PATCH 3/6] Reduce graph transformer PR churn --- gigl/nn/graph_transformer.py | 80 ++++++++++++++++++------------------ 1 file changed, 39 insertions(+), 41 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 1f78f8826..6b2031531 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -29,9 +29,6 @@ heterodata_to_graph_transformer_input, ) -_L2_NORMALIZE_EPS = 1e-6 - - def _get_node_type_positional_encodings( data: torch_geometric.data.hetero_data.HeteroData, node_type: NodeType, @@ -1656,8 +1653,7 @@ def forward( embeddings = self._output_projection(embeddings) if self._should_l2_normalize_embedding_layer_output: - # Default eps=1e-12 underflows to zero in fp16 autocast. - embeddings = F.normalize(embeddings, p=2, dim=-1, eps=_L2_NORMALIZE_EPS) + embeddings = F.normalize(embeddings, p=2, dim=-1) return embeddings @@ -1890,41 +1886,6 @@ 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, @@ -1989,4 +1950,41 @@ def _encode_and_readout( x = self._final_norm(x) x = x * output_valid_mask.unsqueeze(-1).to(x.dtype) - return self._readout_from_encoded_sequences(x=x, valid_mask=valid_mask) + # 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 From 95dac35adf59cd3bff778b53eb27bb0fc899b5ae Mon Sep 17 00:00:00 2001 From: mkolodner Date: Wed, 5 Aug 2026 19:05:18 +0000 Subject: [PATCH 4/6] Align PPR token features with sampler metadata --- gigl/nn/graph_transformer.py | 24 ++-- gigl/transforms/graph_transformer.py | 134 ++++++++++-------- tests/unit/nn/graph_transformer_test.py | 6 +- .../unit/transforms/graph_transformer_test.py | 14 +- 4 files changed, 93 insertions(+), 85 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index 6b2031531..de89c892e 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -22,7 +22,7 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType from gigl.transforms.graph_transformer import ( - PPR_RELATION_FEATURES_NAME, + PPR_FEATURES_NAME, PPR_WEIGHT_FEATURE_NAME, SequenceAuxiliaryData, TokenInputData, @@ -1100,14 +1100,14 @@ 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. 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 + edge weights in PPR mode. The reserved ``"ppr_features"`` name + resolves to additional PPR edge-attr columns after the first weight + column, such as hop or typed-channel metadata. 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``. + Example: ``['hop_distance', 'ppr_weight', 'ppr_features']`` for + continuous features, or ``['hop_distance']`` when ``hop_distance`` + will be embedded via ``anchor_based_input_embedding_dict``. anchor_based_input_embedding_dict: Optional ModuleDict mapping a subset of ``anchor_based_input_attr_names`` to per-attribute embedding layers. These attributes are treated as discrete indices and their @@ -1313,7 +1313,7 @@ def __init__( pairwise_bias_attr_names = pairwise_attention_bias_attr_names or [] ppr_reserved_feature_names = { PPR_WEIGHT_FEATURE_NAME, - PPR_RELATION_FEATURES_NAME, + PPR_FEATURES_NAME, } if ppr_reserved_feature_names & set(pairwise_bias_attr_names): raise ValueError( @@ -1329,9 +1329,9 @@ def __init__( "Reserved PPR anchor-relative features require " "sequence_construction_method='ppr'." ) - if PPR_RELATION_FEATURES_NAME in anchor_bias_attr_names: + if PPR_FEATURES_NAME in anchor_bias_attr_names: raise ValueError( - f"'{PPR_RELATION_FEATURES_NAME}' is a multi-column token-input " + f"'{PPR_FEATURES_NAME}' is a multi-column token-input " "feature and cannot be used as attention bias." ) self._sequence_construction_method = sequence_construction_method @@ -1569,7 +1569,7 @@ def forward( | set(self._pairwise_attention_bias_attr_names or []) ) relative_pe_attr_names.difference_update( - {PPR_WEIGHT_FEATURE_NAME, PPR_RELATION_FEATURES_NAME} + {PPR_WEIGHT_FEATURE_NAME, PPR_FEATURES_NAME} ) if relative_pe_attr_names: for attr_name in sorted(relative_pe_attr_names): diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index 2e520a945..a22cb9d2b 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -79,7 +79,7 @@ class SequenceAuxiliaryData(TypedDict): PPR_WEIGHT_FEATURE_NAME = "ppr_weight" -PPR_RELATION_FEATURES_NAME = "ppr_relation_features" +PPR_FEATURES_NAME = "ppr_features" class _TokenOccurrenceIndex(NamedTuple): @@ -134,7 +134,7 @@ def heterodata_to_graph_transformer_input( ``"ppr"`` uses outgoing ``(anchor_type, "ppr", neighbor_type)`` edges, sorted by descending PPR weight from edge-attr column 0. Multi-column edge attrs can also expose remaining columns as - ``"ppr_relation_features"``. (default: ``"khop"``) + ``"ppr_features"``. (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 @@ -145,10 +145,11 @@ 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. 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']. + to PPR edge weights in PPR sequence mode. The reserved + ``"ppr_features"`` name resolves to any additional PPR edge-attr + columns after the first weight column, such as hop or typed-channel + metadata emitted by the sampler. + Example: ['hop_distance', 'ppr_weight', 'ppr_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']. @@ -223,7 +224,7 @@ def heterodata_to_graph_transformer_input( ppr_reserved_feature_names = { PPR_WEIGHT_FEATURE_NAME, - PPR_RELATION_FEATURES_NAME, + PPR_FEATURES_NAME, } requested_ppr_feature_names = set(anchor_bias_attr_names + anchor_input_attr_names) @@ -253,9 +254,9 @@ def heterodata_to_graph_transformer_input( "sequence_construction_method='ppr'." ) - if PPR_RELATION_FEATURES_NAME in anchor_bias_attr_names: + if PPR_FEATURES_NAME in anchor_bias_attr_names: raise ValueError( - f"'{PPR_RELATION_FEATURES_NAME}' is a multi-column token-input feature " + f"'{PPR_FEATURES_NAME}' is a multi-column token-input feature " "and cannot be used as attention bias." ) @@ -291,7 +292,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 + ppr_feature_sequences: Optional[Tensor] = None if sequence_construction_method == "khop": homo_edge_index = homo_data.edge_index # (2, num_edges) if sampling_direction == "in": @@ -317,7 +318,7 @@ def heterodata_to_graph_transformer_input( node_index_sequences, valid_mask, ppr_weight_sequences, - ppr_relation_feature_sequences, + ppr_feature_sequences, ) = _build_sequence_layout_from_ppr_edges( homo_data=homo_data, anchor_indices=anchor_indices, @@ -329,9 +330,7 @@ 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 - ), + return_ppr_features=(PPR_FEATURES_NAME in anchor_input_attr_names), ) else: raise ValueError( @@ -397,14 +396,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, + ppr_feature_sequences=ppr_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, + ppr_feature_sequences=ppr_feature_sequences, ) return ( @@ -474,7 +473,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], + ppr_feature_sequences: Optional[Tensor], ) -> Optional[Tensor]: if not requested_anchor_attr_names: return None @@ -491,7 +490,7 @@ def _compose_anchor_feature_tensor( 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, + ppr_feature_sequences=ppr_feature_sequences, ) ) @@ -503,7 +502,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], + ppr_feature_sequences: Optional[Tensor], ) -> Optional[TokenInputData]: if not requested_anchor_attr_names: return None @@ -519,7 +518,7 @@ def _compose_anchor_feature_dict( 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, + ppr_feature_sequences=ppr_feature_sequences, ) return feature_dict @@ -530,7 +529,7 @@ def _resolve_anchor_feature_sequence( anchor_relative_feature_sequences: Optional[Tensor], feature_index_by_name: dict[str, int], ppr_weight_sequences: Optional[Tensor], - ppr_relation_feature_sequences: Optional[Tensor], + ppr_feature_sequences: Optional[Tensor], ) -> Tensor: if attr_name == PPR_WEIGHT_FEATURE_NAME: if ppr_weight_sequences is None: @@ -538,12 +537,12 @@ def _resolve_anchor_feature_sequence( f"Requested '{PPR_WEIGHT_FEATURE_NAME}' but it was not computed." ) return ppr_weight_sequences - if attr_name == PPR_RELATION_FEATURES_NAME: - if ppr_relation_feature_sequences is None: + if attr_name == PPR_FEATURES_NAME: + if ppr_feature_sequences is None: raise ValueError( - f"Requested '{PPR_RELATION_FEATURES_NAME}' but it was not computed." + f"Requested '{PPR_FEATURES_NAME}' but it was not computed." ) - return ppr_relation_feature_sequences + return ppr_feature_sequences if anchor_relative_feature_sequences is None: raise ValueError("Anchor-relative features were requested but not computed.") @@ -659,7 +658,7 @@ def _build_sequence_layout_from_ppr_edges( num_nodes: int, device: torch.device, return_edge_weights: bool = False, - return_relation_features: bool = False, + return_ppr_features: bool = False, ) -> tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]: """Build sequences directly from outgoing PPR edges for each anchor. @@ -688,7 +687,7 @@ def _build_sequence_layout_from_ppr_edges( dtype=torch.float, device=device, ) - ppr_relation_feature_sequences = None + ppr_feature_sequences = None if include_anchor_first and max_seq_len > 0: node_index_sequences[:, 0] = anchor_indices @@ -702,7 +701,7 @@ def _build_sequence_layout_from_ppr_edges( node_index_sequences, valid_mask, ppr_weight_sequences, - ppr_relation_feature_sequences, + ppr_feature_sequences, ) if not hasattr(homo_data, "edge_attr") or homo_data.edge_attr is None: @@ -710,28 +709,41 @@ def _build_sequence_layout_from_ppr_edges( "sequence_construction_method='ppr' requires homogeneous edge_attr weights." ) - edge_features = homo_data.edge_attr.float() - if edge_features.dim() == 1: - edge_features = edge_features.unsqueeze(1) - elif edge_features.dim() != 2: + raw_edge_features = homo_data.edge_attr.float() + if raw_edge_features.dim() == 1: + raw_edge_features = raw_edge_features.unsqueeze(1) + elif raw_edge_features.dim() != 2: raise ValueError( - f"PPR edge features must be 1D or 2D, got {tuple(edge_features.shape)}." + "PPR edge features must be 1D or 2D, " + f"got {tuple(raw_edge_features.shape)}." ) - if edge_features.size(1) < 1: + if raw_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: + edge_weights = torch.nan_to_num( + raw_edge_features[:, 0], + nan=0.0, + posinf=1.0, + neginf=0.0, + ).clamp_(min=0.0, max=1.0) + ppr_features = None + if raw_edge_features.size(1) > 1: + ppr_features = torch.nan_to_num( + raw_edge_features[:, 1:], + nan=0.0, + posinf=0.0, + neginf=0.0, + ).clamp_(min=0.0) + if return_ppr_features: + if raw_edge_features.size(1) == 1: raise ValueError( - f"Requested '{PPR_RELATION_FEATURES_NAME}' but PPR edge_attr only " + f"Requested '{PPR_FEATURES_NAME}' but PPR edge_attr only " "contains the scalar weight column." ) - ppr_relation_feature_sequences = torch.zeros( - (batch_size, max_seq_len, edge_features.size(1) - 1), + ppr_feature_sequences = torch.zeros( + (batch_size, max_seq_len, raw_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,), @@ -752,15 +764,13 @@ def _build_sequence_layout_from_ppr_edges( node_index_sequences, valid_mask, ppr_weight_sequences, - ppr_relation_feature_sequences, + ppr_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 - ) + all_ppr_features = ppr_features[keep] if ppr_features is not None else None if include_anchor_first: keep = all_dst_idx != anchor_indices[all_anchor_batch_idx] @@ -769,32 +779,30 @@ def _build_sequence_layout_from_ppr_edges( node_index_sequences, valid_mask, ppr_weight_sequences, - ppr_relation_feature_sequences, + ppr_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 + all_ppr_features = ( + all_ppr_features[keep] if all_ppr_features is not None else 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] - all_relation_features = ( - all_relation_features[weight_order] - if all_relation_features is not None - else None + all_ppr_features = ( + all_ppr_features[weight_order] if all_ppr_features is not None else None ) 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 + sorted_ppr_features = ( + all_ppr_features[batch_order] + if all_ppr_features is not None else None ) @@ -813,9 +821,9 @@ 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 + valid_ppr_features = ( + sorted_ppr_features[valid] + if sorted_ppr_features is not None else None ) @@ -826,18 +834,18 @@ def _build_sequence_layout_from_ppr_edges( valid_weights.float() ) if ( - ppr_relation_feature_sequences is not None - and valid_relation_features is not None + ppr_feature_sequences is not None + and valid_ppr_features is not None ): - ppr_relation_feature_sequences[valid_batch_idx, valid_positions] = ( - valid_relation_features.float() + ppr_feature_sequences[valid_batch_idx, valid_positions] = ( + valid_ppr_features.float() ) return ( node_index_sequences, valid_mask, ppr_weight_sequences, - ppr_relation_feature_sequences, + ppr_feature_sequences, ) diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index e24276eb2..69109cb06 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -15,7 +15,7 @@ GraphTransformerEncoderLayer, ) from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation -from gigl.transforms.graph_transformer import PPR_RELATION_FEATURES_NAME +from gigl.transforms.graph_transformer import PPR_FEATURES_NAME from tests.test_assets.test_case import TestCase @@ -1960,7 +1960,7 @@ 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: + def test_forward_supports_ppr_token_input_features(self) -> None: data = _create_user_graph_with_ppr_edges() data["user", "ppr", "user"].edge_attr = torch.tensor( [ @@ -1974,7 +1974,7 @@ def test_forward_supports_ppr_relation_token_input_features(self) -> None: 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], + anchor_based_input_attr_names=[PPR_FEATURES_NAME], ) encoder.eval() diff --git a/tests/unit/transforms/graph_transformer_test.py b/tests/unit/transforms/graph_transformer_test.py index ec95b278d..11531538a 100644 --- a/tests/unit/transforms/graph_transformer_test.py +++ b/tests/unit/transforms/graph_transformer_test.py @@ -11,7 +11,7 @@ from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.transforms.graph_transformer import ( - PPR_RELATION_FEATURES_NAME, + PPR_FEATURES_NAME, PPR_WEIGHT_FEATURE_NAME, _get_k_hop_neighbors_sparse, heterodata_to_graph_transformer_input, @@ -638,7 +638,7 @@ 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): + def test_ppr_sequence_can_return_ppr_token_input_features(self): data = create_ppr_sequence_hetero_data() data["user", "ppr", "item"].edge_attr = torch.tensor( [ @@ -662,7 +662,7 @@ def test_ppr_sequence_can_return_relation_token_input_features(self): sequence_construction_method="ppr", anchor_based_input_attr_names=[ PPR_WEIGHT_FEATURE_NAME, - PPR_RELATION_FEATURES_NAME, + PPR_FEATURES_NAME, ], ) @@ -671,7 +671,7 @@ def test_ppr_sequence_can_return_relation_token_input_features(self): self.assertEqual( set(token_input.keys()), - {PPR_WEIGHT_FEATURE_NAME, PPR_RELATION_FEATURES_NAME}, + {PPR_WEIGHT_FEATURE_NAME, PPR_FEATURES_NAME}, ) self.assertTrue( torch.allclose( @@ -686,7 +686,7 @@ def test_ppr_sequence_can_return_relation_token_input_features(self): ) self.assertTrue( torch.allclose( - token_input[PPR_RELATION_FEATURES_NAME], + token_input[PPR_FEATURES_NAME], torch.tensor( [ [[0.0, 0.0], [9.0, 90.0], [6.0, 60.0], [4.0, 40.0]], @@ -719,7 +719,7 @@ def test_ppr_sequence_sorts_multi_column_edge_attr_by_weight_column(self): anchor_node_type="user", sequence_construction_method="ppr", anchor_based_attention_bias_attr_names=[PPR_WEIGHT_FEATURE_NAME], - anchor_based_input_attr_names=[PPR_RELATION_FEATURES_NAME], + anchor_based_input_attr_names=[PPR_FEATURES_NAME], ) anchor_bias = sequence_auxiliary_data["anchor_bias"] @@ -741,7 +741,7 @@ def test_ppr_sequence_sorts_multi_column_edge_attr_by_weight_column(self): ) self.assertTrue( torch.allclose( - token_input[PPR_RELATION_FEATURES_NAME], + token_input[PPR_FEATURES_NAME], torch.tensor([[[0.0, 0.0], [0.0, 1.0], [1.0, 0.0]]]), ) ) From d3f33327a48cd2eba4aa971ebf4e68a1ff86856c Mon Sep 17 00:00:00 2001 From: mkolodner Date: Fri, 14 Aug 2026 19:13:49 +0000 Subject: [PATCH 5/6] Support PPR sequences with preserved original edges --- gigl/transforms/graph_transformer.py | 44 +++++++++------- .../unit/transforms/graph_transformer_test.py | 51 ++++++++++++++++++- 2 files changed, 74 insertions(+), 21 deletions(-) diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index a22cb9d2b..91e6cd438 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -62,6 +62,7 @@ import torch from torch import Tensor from torch_geometric.data import Data, HeteroData +from torch_geometric.typing import EdgeType as PyGEdgeType from torch_geometric.typing import NodeType from torch_geometric.utils import to_torch_sparse_tensor @@ -80,6 +81,7 @@ class SequenceAuxiliaryData(TypedDict): PPR_WEIGHT_FEATURE_NAME = "ppr_weight" PPR_FEATURES_NAME = "ppr_features" +PPR_RELATION_NAME = "ppr" class _TokenOccurrenceIndex(NamedTuple): @@ -133,6 +135,8 @@ def heterodata_to_graph_transformer_input( ``"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 column 0. + Non-PPR edges can remain on the batch for relation-aware attention + and message features, but they are ignored for PPR sequence layout. Multi-column edge attrs can also expose remaining columns as ``"ppr_features"``. (default: ``"khop"``) include_anchor_first: If True, anchor node is always first in sequence. @@ -260,13 +264,19 @@ def heterodata_to_graph_transformer_input( "and cannot be used as attention bias." ) + ppr_edge_types: list[PyGEdgeType] = [] if sequence_construction_method == "ppr": - _validate_ppr_sequence_input(data) + ppr_edge_types = _validate_ppr_sequence_input(data) device = data[anchor_node_type].x.device # Convert to homogeneous for easier neighborhood extraction - homo_data = data.to_homogeneous() + sequence_data = ( + data.edge_type_subgraph(ppr_edge_types) + if sequence_construction_method == "ppr" + else data + ) + homo_data = sequence_data.to_homogeneous() homo_x = homo_data.x # (total_nodes, feature_dim) num_nodes = homo_data.num_nodes @@ -431,25 +441,29 @@ def _get_node_type_offsets( return offsets -def _validate_ppr_sequence_input(data: HeteroData) -> None: +def _validate_ppr_sequence_input(data: HeteroData) -> list[PyGEdgeType]: if not data.edge_types: raise ValueError( "sequence_construction_method='ppr' requires at least one PPR edge type." ) - if any(edge_type[1] != "ppr" for edge_type in data.edge_types): + ppr_edge_types = [ + edge_type for edge_type in data.edge_types if edge_type[1] == PPR_RELATION_NAME + ] + if not ppr_edge_types: raise ValueError( - "sequence_construction_method='ppr' expects the hetero batch to contain " - f"only PPR edges, got edge types: {data.edge_types}." + "sequence_construction_method='ppr' requires at least one PPR edge type, " + f"got edge types: {data.edge_types}." ) - for edge_type in data.edge_types: + for edge_type in ppr_edge_types: edge_store = data[edge_type] if not hasattr(edge_store, "edge_attr") or edge_store.edge_attr is None: raise ValueError( "sequence_construction_method='ppr' requires every PPR edge type to " f"have edge_attr weights, but {edge_type} is missing them." ) + return ppr_edge_types def _get_sparse_feature_matrices( @@ -714,8 +728,7 @@ def _build_sequence_layout_from_ppr_edges( raw_edge_features = raw_edge_features.unsqueeze(1) elif raw_edge_features.dim() != 2: raise ValueError( - "PPR edge features must be 1D or 2D, " - f"got {tuple(raw_edge_features.shape)}." + f"PPR edge features must be 1D or 2D, got {tuple(raw_edge_features.shape)}." ) if raw_edge_features.size(1) < 1: raise ValueError("PPR edge features must contain at least one weight column.") @@ -801,9 +814,7 @@ def _build_sequence_layout_from_ppr_edges( sorted_dst_idx = all_dst_idx[batch_order] sorted_weights = all_weights[batch_order] sorted_ppr_features = ( - all_ppr_features[batch_order] - if all_ppr_features is not None - else None + all_ppr_features[batch_order] if all_ppr_features is not None else None ) n = sorted_batch_idx.size(0) @@ -822,9 +833,7 @@ def _build_sequence_layout_from_ppr_edges( valid_dst_idx = sorted_dst_idx[valid] valid_weights = sorted_weights[valid] valid_ppr_features = ( - sorted_ppr_features[valid] - if sorted_ppr_features is not None - else None + sorted_ppr_features[valid] if sorted_ppr_features is not None else None ) node_index_sequences[valid_batch_idx, valid_positions] = valid_dst_idx @@ -833,10 +842,7 @@ def _build_sequence_layout_from_ppr_edges( ppr_weight_sequences[valid_batch_idx, valid_positions, 0] = ( valid_weights.float() ) - if ( - ppr_feature_sequences is not None - and valid_ppr_features is not None - ): + if ppr_feature_sequences is not None and valid_ppr_features is not None: ppr_feature_sequences[valid_batch_idx, valid_positions] = ( valid_ppr_features.float() ) diff --git a/tests/unit/transforms/graph_transformer_test.py b/tests/unit/transforms/graph_transformer_test.py index 11531538a..2d4890d7b 100644 --- a/tests/unit/transforms/graph_transformer_test.py +++ b/tests/unit/transforms/graph_transformer_test.py @@ -571,10 +571,10 @@ def test_ppr_sequence_construction_sorts_tokens_by_weight(self): torch.equal(valid_mask[1], torch.tensor([True, True, True, False])) ) - def test_ppr_sequence_construction_requires_only_ppr_relations(self): + def test_ppr_sequence_construction_requires_ppr_relations(self): data = create_simple_hetero_data() - with self.assertRaisesRegex(ValueError, "contain only PPR edges"): + with self.assertRaisesRegex(ValueError, "requires at least one PPR edge type"): heterodata_to_graph_transformer_input( data=data, batch_size=1, @@ -583,6 +583,53 @@ def test_ppr_sequence_construction_requires_only_ppr_relations(self): sequence_construction_method="ppr", ) + def test_ppr_sequence_ignores_original_edges_but_keeps_relation_indices(self): + data = create_ppr_sequence_hetero_data() + user = NodeType("user") + item = NodeType("item") + buys = EdgeType(user, Relation("buys"), item) + data[buys.tuple_repr()].edge_index = torch.tensor([[0], [0]]) + data[buys.tuple_repr()].edge_attr = torch.tensor([[99.0, 99.0, 99.0]]) + + sequences, valid_mask, 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", + relation_edge_types=[buys], + ) + ) + + self.assertTrue( + torch.allclose( + sequences[0], + torch.tensor( + [ + [10.0, 0.0], + [0.0, 21.0], + [0.0, 20.0], + [11.0, 0.0], + ] + ), + ) + ) + self.assertTrue( + torch.equal(valid_mask[1], torch.tensor([True, True, True, False])) + ) + + pairwise_relation_indices = sequence_auxiliary_data["pairwise_relation_indices"] + self.assertIsNotNone(pairwise_relation_indices) + assert pairwise_relation_indices is not None + self.assertEqual( + {tuple(coord) for coord in pairwise_relation_indices.tolist()}, + { + (0, 2, 0, 0), + (1, 1, 2, 0), + }, + ) + def test_ppr_sequence_can_return_token_input_and_attention_bias_features(self): data = create_ppr_sequence_hetero_data() From 423ff50ec96eb61e4c319e8db191a831f50f90df Mon Sep 17 00:00:00 2001 From: mkolodner Date: Fri, 14 Aug 2026 21:01:05 +0000 Subject: [PATCH 6/6] Clarify PPR reserved feature validation --- gigl/nn/graph_transformer.py | 23 ++++++++++++++--------- gigl/transforms/graph_transformer.py | 22 +++++++++++++--------- 2 files changed, 27 insertions(+), 18 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index de89c892e..cae8e8b05 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -29,6 +29,7 @@ heterodata_to_graph_transformer_input, ) + def _get_node_type_positional_encodings( data: torch_geometric.data.hetero_data.HeteroData, node_type: NodeType, @@ -1315,18 +1316,22 @@ def __init__( PPR_WEIGHT_FEATURE_NAME, PPR_FEATURES_NAME, } - if ppr_reserved_feature_names & set(pairwise_bias_attr_names): + requested_pairwise_ppr_feature_names = ppr_reserved_feature_names & set( + pairwise_bias_attr_names + ) + requested_anchor_ppr_feature_names = ppr_reserved_feature_names & set( + anchor_bias_attr_names + anchor_input_attr_names + ) + if requested_pairwise_ppr_feature_names: raise ValueError( - f"{sorted(ppr_reserved_feature_names)} are anchor-relative features and " - "cannot be used as pairwise attention bias." + "PPR reserved features " + f"{sorted(requested_pairwise_ppr_feature_names)} are anchor-relative " + "features and cannot be used as pairwise attention bias." ) - if ( - ppr_reserved_feature_names - & set(anchor_bias_attr_names + anchor_input_attr_names) - and sequence_construction_method != "ppr" - ): + if requested_anchor_ppr_feature_names and sequence_construction_method != "ppr": raise ValueError( - "Reserved PPR anchor-relative features require " + "PPR reserved features " + f"{sorted(requested_anchor_ppr_feature_names)} require " "sequence_construction_method='ppr'." ) if PPR_FEATURES_NAME in anchor_bias_attr_names: diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index 91e6cd438..e5ee236f8 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -230,12 +230,18 @@ def heterodata_to_graph_transformer_input( PPR_WEIGHT_FEATURE_NAME, PPR_FEATURES_NAME, } - requested_ppr_feature_names = set(anchor_bias_attr_names + anchor_input_attr_names) + requested_pairwise_ppr_feature_names = ppr_reserved_feature_names & set( + pairwise_bias_attr_names + ) + requested_anchor_ppr_feature_names = ppr_reserved_feature_names & set( + anchor_bias_attr_names + anchor_input_attr_names + ) - if ppr_reserved_feature_names & set(pairwise_bias_attr_names): + if requested_pairwise_ppr_feature_names: raise ValueError( - f"{sorted(ppr_reserved_feature_names)} are anchor-relative features and cannot " - "be used as pairwise attention bias." + "PPR reserved features " + f"{sorted(requested_pairwise_ppr_feature_names)} are anchor-relative " + "features and cannot be used as pairwise attention bias." ) if sampling_direction not in {"in", "out"}: @@ -249,12 +255,10 @@ def heterodata_to_graph_transformer_input( "sequence_construction_method='ppr' supports only sampling_direction='out'." ) - if ( - ppr_reserved_feature_names & requested_ppr_feature_names - and sequence_construction_method != "ppr" - ): + if requested_anchor_ppr_feature_names and sequence_construction_method != "ppr": raise ValueError( - "Reserved PPR anchor-relative features require " + "PPR reserved features " + f"{sorted(requested_anchor_ppr_feature_names)} require " "sequence_construction_method='ppr'." )