Add PPR relation token features to graph transformer

[mkolodner-sc]

Summary

Adds Graph Transformer support for generic PPR edge-attribute metadata.

PPR can emit multi-column edge_attr tensors where column 0 is the scalar PPR weight and the remaining columns contain sampler metadata. This PR lets Graph Transformer consume those extra columns as token-input features while preserving the existing scalar PPR path.

This also works with the current GiGL main PPR output shape, where heterogeneous PPR batches may include preserved original message-passing edges alongside virtual PPR edges. Graph Transformer uses only virtual (anchor_type, "ppr", neighbor_type) edges for PPR token sequence construction, while still allowing preserved original edges to participate in relation-aware attention and value/message edge-type features.

What This Enables

A PPR edge can now carry:

• ppr_weight: the first edge_attr column, used as the scalar PPR score.
• ppr_features: all remaining edge_attr columns, projected into the token embedding as continuous token features.

For current GiGL PPR outputs:

• Regular PPR emits edge_attr as [ppr_score, hop_proximity].
• Typed-channel PPR emits edge_attr as [best_score, hop_proximity, (channel_score, channel_hop_proximity, channel_presence), ...].

hop_proximity is 1 / (1 + hop), so anchor/self is 1.0, one-hop is 0.5, two-hop is about 0.333, and so on. For channel-level hop proximity, missing channels use 0.0 and should be interpreted using the corresponding channel presence bit.

This allows PPR metadata such as global hop proximity, channel scores, channel hop proximity, and channel presence bits to flow into Graph Transformer tokens without changing how scalar PPR batches behave.

Changes

• Adds reserved feature name ppr_features.
• Supports scalar and multi-column PPR edge_attr.
• Keeps ppr_weight mapped to edge_attr[..., 0].
• Maps ppr_features to edge_attr[..., 1:].
• Preserves metadata values in ppr_features.
• Keeps ppr_weight clamped as a scalar score in [0, 1].
• Sorts PPR tokens by descending edge_attr[..., 0] and carries feature columns through the same permutation.
• Validates that ppr_features is used only as token input, not attention bias.
• Updates token-input handling to support multi-column continuous features.
• Supports PPR batches that include preserved original message-passing edges.
• Uses only virtual PPR edges for PPR sequence construction.
• Keeps preserved original edges available for relation-aware attention and value/message edge-type features.
• Clarifies PPR reserved-feature validation by naming the requested PPR feature intersections explicitly.
• Adds transform and encoder tests for PPR feature extraction, ordering, mixed PPR/original-edge batches, and forward-pass support.

Toy Example

Suppose typed-channel PPR emits virtual PPR edges from an anchor to selected tokens.

Each PPR edge represents one selected (anchor, token) pair. The edge_attr row describes how PPR reached/scored that token:

• Column 0 is the scalar best PPR score.
• Column 1 is the global hop proximity.
• Each typed channel contributes:
  • channel score
  • channel hop proximity
  • channel presence bit

For example, with two typed PPR channels:

data["user", "ppr", "item"].edge_attr = torch.tensor([
    # best_ppr, global_hop_prox, ch0_score, ch0_hop_prox, ch0_present, ch1_score, ch1_hop_prox, ch1_present
    [0.90,     0.50,            0.90,      0.50,         1.0,         0.00,      0.00,         0.0],
    [0.75,     0.33,            0.10,      0.33,         1.0,         0.75,      0.33,         1.0],
    [0.40,     0.20,            0.00,      0.00,         0.0,         0.40,      0.20,         1.0],
])

Graph Transformer interprets this as:

ppr_weight = edge_attr[:, 0:1]
ppr_features = edge_attr[:, 1:]

A model can consume both as token-input features:

encoder = GraphTransformerEncoder(
    ...,
    sequence_construction_method="ppr",
    anchor_based_input_attr_names=[
        "ppr_weight",
        "ppr_features",
    ],
)

These continuous features are concatenated, passed through the learned token-input projection, and added to the corresponding token embeddings before the Transformer layers run.

Preserved Original Edges

Current heterogeneous PPR can preserve original graph edges alongside virtual PPR edges when the sampler is configured to include sampled edges.

In that mixed batch shape:

• Virtual PPR edges drive PPR token sequence construction.
• Original message-passing edges are ignored for PPR token ordering.
• Original message-passing edges remain available for relation-aware attention and value/message edge-type features.
• Original edges only contribute when both endpoints are present in the PPR-selected token sequence.

This enables PPR token metadata and original-edge relation/value modeling to be used together.

Example:

encoder = GraphTransformerEncoder(
    ...,
    sequence_construction_method="ppr",
    anchor_based_input_attr_names=[
        "ppr_weight",
        "ppr_features",
    ],
    edge_type_to_feat_dim_map={
        user_to_item: 0,
        item_to_user: 0,
    },
    relation_attention_mode="edge_type_bilinear",
    relation_message_mode="edge_type_linear",
)

Validation

• Added focused unit coverage for multi-column PPR edge attrs and Graph Transformer token-input consumption.
• Added coverage for mixed virtual-PPR plus preserved-original-edge batches.
• Added encoder forward-pass coverage for PPR token-input metadata.
• Verified stale ppr_relation_features references were removed.
• Verified ppr_features preserves metadata columns rather than clamping hop-like columns.
• Ran syntax, whitespace, ruff, and focused runtime checks on the changed Graph Transformer files and tests.