Add PPR relation token features to graph transformer

gigl/nn/graph_transformer.py

@@ -22,6 +22,7 @@
 
 from gigl.src.common.types.graph_data import EdgeType, NodeType
 from gigl.transforms.graph_transformer import (
+    PPR_FEATURES_NAME,
     PPR_WEIGHT_FEATURE_NAME,
     SequenceAuxiliaryData,
     TokenInputData,
@@ -1100,11 +1101,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. 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``.
+            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', '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
@@ -1308,19 +1312,33 @@ 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_FEATURES_NAME,
+        }
+        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"'{PPR_WEIGHT_FEATURE_NAME}' is an anchor-relative feature 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_WEIGHT_FEATURE_NAME in 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(
-                "The reserved anchor-relative feature 'ppr_weight' requires "
+                "PPR reserved features "
+                f"{sorted(requested_anchor_ppr_feature_names)} require "
                 "sequence_construction_method='ppr'."
             )
+        if PPR_FEATURES_NAME in anchor_bias_attr_names:
+            raise ValueError(
+                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
         self._sampling_direction = sampling_direction
         self._sequence_positional_encoding_type = sequence_positional_encoding_type
@@ -1385,7 +1403,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 +1568,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_FEATURES_NAME}
+        )
         if relative_pe_attr_names:
             for attr_name in sorted(relative_pe_attr_names):
                 if hasattr(data, attr_name):
@@ -1724,11 +1741,11 @@ 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
+            )
             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