"""
Transform HeteroData to Graph Transformer sequence input.
This module provides functionality to convert PyG HeteroData objects (typically
batched 2-hop subgraphs) into sequence format suitable for Graph Transformers.
For each anchor node in the batch, the transform extracts its k-hop neighborhood
and creates a fixed-length sequence of node features with padding.
Example Usage:
>>> from torch_geometric.data import HeteroData
>>> from gigl.transforms.graph_transformer import heterodata_to_graph_transformer_input
>>>
>>> # Create batched HeteroData (e.g., from NeighborLoader)
>>> # First batch_size nodes in each node type are anchor nodes
>>> data = HeteroData()
>>> data['user'].x = torch.randn(100, 64) # 100 users, first N are anchors
>>> data['item'].x = torch.randn(50, 32)
>>> data['user', 'buys', 'item'].edge_index = ...
>>>
>>> # Transform to Graph Transformer input
>>> sequences, valid_mask, attention_bias_data = heterodata_to_graph_transformer_input(
... data=data,
... batch_size=32,
... max_seq_len=128,
... anchor_node_type='user',
... )
>>> # sequences: (batch_size, max_seq_len, feature_dim)
>>> # valid_mask: (batch_size, max_seq_len)
``tensors_to_graph_transformer_input`` takes the same arguments with the
batch's ``x_dict`` / ``edge_index_dict`` (and optional ``edge_attr_dict``,
``graph_attr_dict``) instead of a ``HeteroData``.
With Relative Encodings:
Relative encodings stored as sparse graph-level attributes can be returned as
raw attention-bias features:
>>> from torch_geometric.transforms import Compose
>>> from gigl.transforms.add_positional_encodings import (
... AddHeteroRandomWalkEncodings,
... AddHeteroHopDistanceEncoding,
... )
>>>
>>> # First apply PE transforms to the data
>>> pe_transform = Compose([
... AddHeteroRandomWalkEncodings(walk_length=8),
... AddHeteroHopDistanceEncoding(h_max=5),
... ])
>>> data = pe_transform(data)
>>>
>>> # Transform to sequences with relative bias features
>>> sequences, valid_mask, attention_bias_data = heterodata_to_graph_transformer_input(
... data=data,
... batch_size=32,
... max_seq_len=128,
... anchor_node_type='user',
... anchor_based_attention_bias_attr_names=['hop_distance'],
... )
>>> # sequences: (batch_size, max_seq_len, feature_dim)
>>> # attention_bias_data['anchor_bias']: (batch_size, max_seq_len, 1)
"""
from typing import Literal, NamedTuple, Optional, TypedDict
import torch
from jaxtyping import Bool, Float, Int64
from torch import Tensor
from torch_geometric.data import HeteroData
from torch_geometric.typing import EdgeType, NodeType
from torch_geometric.utils import to_torch_sparse_tensor
from gigl.src.common.types.graph_data import EdgeType as GiGLEdgeType
[docs]
class SequenceAuxiliaryData(TypedDict):
[docs]
anchor_bias: Optional[Float[Tensor, "anchors sequence"]]
[docs]
pairwise_bias: Optional[Float[Tensor, "anchors sequence sequence"]]
[docs]
pairwise_relation_indices: Optional[Int64[Tensor, "relation_edges 4"]]
[docs]
pairwise_nonmissing_indices: Optional[Int64[Tensor, "nonmissing_edges 3"]]
[docs]
PPR_WEIGHT_FEATURE_NAME = "ppr_weight"
[docs]
PPR_RELATION_FEATURES_NAME = "ppr_relation_features"
[docs]
PPR_EXPANDED_FEATURE_NAME = "ppr_is_expanded"
def _validate_reserved_anchor_feature_usage(
pairwise_bias_attr_names: list[str],
anchor_bias_attr_names: list[str],
anchor_input_attr_names: list[str],
sequence_construction_method: str,
) -> None:
"""Validate placement of the reserved PPR anchor-relative features.
``ppr_weight`` and ``ppr_relation_features`` are produced per (anchor,
neighbor) edge, so they may only be requested as anchor-relative token
inputs (and, for the scalar ``ppr_weight``, as a scalar bias) under
``sequence_construction_method='ppr'``. Shared by
``heterodata_to_graph_transformer_input`` and ``GraphTransformerEncoder`` so
the rules stay defined in one place.
"""
# Neither reserved PPR feature may be used as a pairwise attention bias;
# both are anchor-relative (produced per (anchor, neighbor) edge).
for name in (PPR_WEIGHT_FEATURE_NAME, PPR_RELATION_FEATURES_NAME):
if name in pairwise_bias_attr_names:
raise ValueError(
f"'{name}' is an anchor-relative feature and cannot be used as "
"pairwise attention bias."
)
# ppr_relation_features is multi-column, so it can only be a token input,
# never a scalar attention bias.
if PPR_RELATION_FEATURES_NAME in anchor_bias_attr_names:
raise ValueError(
f"'{PPR_RELATION_FEATURES_NAME}' is multi-column token input and "
"cannot be used as scalar attention bias."
)
# Either reserved feature, used as any anchor feature (scalar bias or token
# input), requires the PPR sequence path.
for name in (PPR_WEIGHT_FEATURE_NAME, PPR_RELATION_FEATURES_NAME):
if (
name in anchor_bias_attr_names + anchor_input_attr_names
and sequence_construction_method != "ppr"
):
raise ValueError(
f"The reserved anchor-relative feature '{name}' requires "
"sequence_construction_method='ppr'."
)
class _TokenOccurrenceIndex(NamedTuple):
batch_indices: Tensor
positions: Tensor
node_indices: Tensor
sorted_node_indices: Tensor
node_sort_perm: Tensor
sorted_batch_node_keys: Tensor
batch_node_sort_perm: Tensor
def _get_node_type_offsets(x_dict: dict[NodeType, Tensor]) -> dict[NodeType, int]:
offsets: dict[NodeType, int] = {}
offset = 0
for node_type, x in x_dict.items():
offsets[node_type] = offset
offset += x.size(0)
return offsets
def _concat_tensors(parts: list[Tensor], dim: int) -> Tensor:
"""``torch.cat`` that returns a single part as is, without a copy."""
return parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim)
def _to_homogeneous_edge_index(
edge_index_dict: dict[EdgeType, Tensor],
node_type_offsets: dict[NodeType, int],
device: torch.device,
) -> Tensor:
"""Concatenates edge indices the way ``HeteroData.to_homogeneous`` does.
Each edge type's ids are shifted by its node type offsets, and edge types
keep ``edge_index_dict``'s order.
Returns:
``[2, num_edges]`` edge index; ``[2, 0]`` if there are no edges.
"""
edge_index_parts: list[Tensor] = []
for (src_type, _, dst_type), edge_index in edge_index_dict.items():
src_offset = node_type_offsets[src_type]
dst_offset = node_type_offsets[dst_type]
if src_offset or dst_offset:
edge_index = torch.stack(
[edge_index[0] + src_offset, edge_index[1] + dst_offset]
)
edge_index_parts.append(edge_index)
if not edge_index_parts:
return torch.zeros((2, 0), dtype=torch.long, device=device)
return _concat_tensors(edge_index_parts, dim=1)
def _validate_ppr_sequence_input(
edge_index_dict: dict[EdgeType, Tensor],
edge_attr_dict: dict[EdgeType, Tensor],
) -> None:
# The batch may also contain original relation edge types when
# include_sampled_edges is on (they feed the relation-message channel via
# _lookup_pairwise_relation_indices); only the "ppr" edges drive the sequence.
ppr_edge_types = [et for et in edge_index_dict if et[1] == "ppr"]
if not ppr_edge_types:
raise ValueError(
"sequence_construction_method='ppr' requires at least one PPR edge type."
)
for edge_type in ppr_edge_types:
if edge_type not in edge_attr_dict:
raise ValueError(
"sequence_construction_method='ppr' requires every PPR edge type to "
f"have edge_attr weights, but {edge_type} is missing them."
)
def _get_sparse_feature_matrices(
graph_attr_dict: dict[str, Tensor],
attr_names: Optional[list[str]],
missing_attr_error_prefix: str,
) -> list[Tensor]:
matrices: list[Tensor] = []
for attr_name in attr_names or []:
if attr_name not in graph_attr_dict:
raise ValueError(
f"{missing_attr_error_prefix} '{attr_name}' not found in data. "
"Make sure to apply the corresponding transform first."
)
matrices.append(graph_attr_dict[attr_name])
return matrices
def _compose_anchor_feature_tensor(
anchor_relative_feature_sequences: Optional[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
feature_parts: list[Tensor] = []
feature_index_by_name = {
attr_name: idx for idx, attr_name in enumerate(available_anchor_attr_names)
}
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 attr_name == PPR_RELATION_FEATURES_NAME:
if ppr_relation_feature_sequences is None:
raise ValueError(
f"Requested '{PPR_RELATION_FEATURES_NAME}' but it was not computed."
)
feature_parts.append(ppr_relation_feature_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]
)
return torch.cat(feature_parts, dim=-1)
def _compose_anchor_feature_dict(
anchor_relative_feature_sequences: Optional[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[TokenInputData]:
if not requested_anchor_attr_names:
return None
feature_dict: TokenInputData = {}
feature_index_by_name = {
attr_name: idx for idx, attr_name in enumerate(available_anchor_attr_names)
}
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
if attr_name == PPR_RELATION_FEATURES_NAME:
if ppr_relation_feature_sequences is None:
raise ValueError(
f"Requested '{PPR_RELATION_FEATURES_NAME}' but it was not computed."
)
feature_dict[attr_name] = ppr_relation_feature_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_dict[attr_name] = anchor_relative_feature_sequences[
..., feature_idx : feature_idx + 1
]
return feature_dict
def _build_sequence_layout_from_sparse_neighbors(
reachable: Tensor,
anchor_indices: Tensor,
max_seq_len: int,
include_anchor_first: bool,
device: torch.device,
) -> tuple[Tensor, Tensor]:
"""
Build padded node-index sequences and a validity mask from reachability data.
Returns:
node_index_sequences: (batch_size, max_seq_len) long tensor with -1 padding
valid_mask: (batch_size, max_seq_len) bool tensor
"""
batch_size = anchor_indices.size(0)
node_index_sequences = torch.full(
(batch_size, max_seq_len),
fill_value=-1,
dtype=torch.long,
device=device,
)
valid_mask = torch.zeros(
(batch_size, max_seq_len),
dtype=torch.bool,
device=device,
)
indices = reachable.indices()
batch_idx = indices[0]
node_idx = indices[1]
if include_anchor_first and max_seq_len > 0:
node_index_sequences[:, 0] = anchor_indices
valid_mask[:, 0] = True
if batch_idx.numel() > 0:
keep = node_idx != anchor_indices[batch_idx]
batch_idx = batch_idx[keep]
node_idx = node_idx[keep]
start_pos = 1
else:
start_pos = 0
if batch_idx.numel() == 0 or start_pos >= max_seq_len:
return node_index_sequences, valid_mask
# Compute within-anchor sequence positions for each neighbor node.
#
# After extracting from the sparse tensor, we have flattened arrays:
# batch_idx = [0, 0, 0, 1, 1, 2, 2, 2, 2] <- which anchor each node belongs to
# node_idx = [5, 7, 9, 3, 8, 1, 4, 6, 10] <- the reachable neighbor node IDs
#
# We need to compute sequence positions (1, 2, 3, ...) for each anchor's neighbors:
# positions = [1, 2, 3, 1, 2, 1, 2, 3, 4]
# ^anchor0^ ^a1^ ^--anchor2--^
#
# This allows scattering into the 2D output:
# node_index_sequences[anchor_idx, position] = node_idx
n = batch_idx.size(0)
# Step 1: Mark where each anchor's neighbor group starts
# is_group_start = [1, 0, 0, 1, 0, 1, 0, 0, 0]
# ^ ^ ^
# anchor0 anchor1 anchor2 starts here
is_group_start = torch.zeros(n, dtype=torch.long, device=device)
is_group_start[0] = 1
if n > 1:
is_group_start[1:] = (batch_idx[1:] != batch_idx[:-1]).long()
# Step 2: Assign each node to its anchor group (0-indexed)
# group_id = [0, 0, 0, 1, 1, 2, 2, 2, 2]
group_id = is_group_start.cumsum(0) - 1
# Step 3: Find the starting index of each group in the flattened array
# group_starts = [0, 3, 5] (indices where each anchor's neighbors begin)
group_starts = torch.nonzero(is_group_start, as_tuple=True)[0]
# Step 4: Compute position = (global_index - group_start) + start_pos
# This gives within-group position offset by start_pos (usually 1 for anchor at pos 0)
# positions = [1, 2, 3, 1, 2, 1, 2, 3, 4]
positions = torch.arange(n, device=device) - group_starts[group_id] + start_pos
# Step 5: Filter out positions that exceed max_seq_len (truncation)
valid = positions < max_seq_len
valid_batch_idx = batch_idx[valid]
valid_positions = positions[valid]
valid_node_idx = node_idx[valid]
# Step 6: Scatter valid nodes into the output tensors
node_index_sequences[valid_batch_idx, valid_positions] = valid_node_idx
valid_mask[valid_batch_idx, valid_positions] = True
return node_index_sequences, valid_mask
def _build_hop_ordered_sequence_layout_from_sparse_edges(
anchor_indices: Tensor,
edge_index: Tensor,
num_nodes: int,
hop_distance: int,
max_seq_len: int,
include_anchor_first: bool,
device: torch.device,
) -> tuple[Tensor, Tensor]:
"""
Build sequences by visiting lower-hop neighbors before higher-hop neighbors.
The implementation keeps only the current sparse frontier and a sorted set
of visited ``(batch, node)`` keys. It avoids a dense
``batch_size x num_nodes`` distance matrix and stops once all sequences are
full.
"""
batch_size = anchor_indices.size(0)
node_index_sequences = torch.full(
(batch_size, max_seq_len),
fill_value=-1,
dtype=torch.long,
device=device,
)
valid_mask = torch.zeros(
(batch_size, max_seq_len),
dtype=torch.bool,
device=device,
)
if max_seq_len == 0:
return node_index_sequences, valid_mask
next_positions = torch.zeros(batch_size, dtype=torch.long, device=device)
if include_anchor_first and max_seq_len > 0:
node_index_sequences[:, 0] = anchor_indices
valid_mask[:, 0] = True
next_positions.fill_(1)
if hop_distance <= 0 or (next_positions >= max_seq_len).all().item():
return node_index_sequences, valid_mask
visited_batch_node_keys = (
torch.arange(batch_size, device=device, dtype=torch.long) * num_nodes
+ anchor_indices
).sort()[0]
current_frontier = torch.sparse_coo_tensor(
torch.stack(
[
torch.arange(batch_size, device=device),
anchor_indices,
]
),
torch.ones(batch_size, dtype=torch.float, device=device),
size=(batch_size, num_nodes),
is_coalesced=True,
)
adj = _binarize_sparse_tensor(
sparse_tensor=to_torch_sparse_tensor(
edge_index,
size=(num_nodes, num_nodes),
),
device=device,
)
for _ in range(hop_distance):
active_batch_mask = next_positions < max_seq_len
current_frontier = _filter_sparse_frontier_by_batch(
frontier=current_frontier,
active_batch_mask=active_batch_mask,
device=device,
)
if current_frontier._nnz() == 0:
break
current_frontier = torch.sparse.mm(current_frontier, adj)
if current_frontier._nnz() == 0:
break
current_frontier = _binarize_sparse_tensor(
sparse_tensor=current_frontier,
device=device,
)
visited_batch_node_keys, current_frontier = (
_append_frontier_tokens_to_sequences(
frontier=current_frontier,
visited_batch_node_keys=visited_batch_node_keys,
next_positions=next_positions,
node_index_sequences=node_index_sequences,
valid_mask=valid_mask,
num_nodes=num_nodes,
device=device,
)
)
if (next_positions >= max_seq_len).all().item():
break
return node_index_sequences, valid_mask
def _empty_sparse_tensor_like(
sparse_tensor: Tensor,
device: torch.device,
) -> Tensor:
return torch.sparse_coo_tensor(
torch.zeros((2, 0), dtype=torch.long, device=device),
torch.zeros(0, dtype=torch.float, device=device),
size=sparse_tensor.shape,
is_coalesced=True,
)
def _binarize_sparse_tensor(
sparse_tensor: Tensor,
device: torch.device,
) -> Tensor:
sparse_tensor = sparse_tensor.coalesce()
return torch.sparse_coo_tensor(
sparse_tensor.indices(),
torch.ones(sparse_tensor._nnz(), device=device, dtype=torch.float),
size=sparse_tensor.shape,
is_coalesced=True,
)
def _filter_sparse_frontier_by_batch(
frontier: Tensor,
active_batch_mask: Tensor,
device: torch.device,
) -> Tensor:
if not frontier.is_coalesced():
frontier = frontier.coalesce()
indices = frontier.indices()
if indices.size(1) == 0:
return frontier
keep = active_batch_mask[indices[0]]
if not keep.any().item():
return _empty_sparse_tensor_like(sparse_tensor=frontier, device=device)
kept_indices = indices[:, keep]
return torch.sparse_coo_tensor(
kept_indices,
torch.ones(kept_indices.size(1), dtype=torch.float, device=device),
size=frontier.shape,
is_coalesced=True,
)
def _append_frontier_tokens_to_sequences(
frontier: Tensor,
visited_batch_node_keys: Tensor,
next_positions: Tensor,
node_index_sequences: Tensor,
valid_mask: Tensor,
num_nodes: int,
device: torch.device,
) -> tuple[Tensor, Tensor]:
if not frontier.is_coalesced():
frontier = frontier.coalesce()
indices = frontier.indices()
if indices.size(1) == 0:
return visited_batch_node_keys, frontier
batch_idx = indices[0]
node_idx = indices[1]
batch_node_keys = batch_idx.long() * num_nodes + node_idx.long()
if visited_batch_node_keys.numel() > 0:
insert_pos = torch.searchsorted(visited_batch_node_keys, batch_node_keys)
insert_pos_clamped = insert_pos.clamp(max=visited_batch_node_keys.size(0) - 1)
is_visited = visited_batch_node_keys[insert_pos_clamped] == batch_node_keys
is_new = ~is_visited
else:
is_new = torch.ones_like(batch_node_keys, dtype=torch.bool)
if not is_new.any().item():
return visited_batch_node_keys, _empty_sparse_tensor_like(
sparse_tensor=frontier,
device=device,
)
new_batch_idx = batch_idx[is_new]
new_node_idx = node_idx[is_new]
new_batch_node_keys = batch_node_keys[is_new]
has_space = next_positions[new_batch_idx] < node_index_sequences.size(1)
if not has_space.any().item():
return visited_batch_node_keys, _empty_sparse_tensor_like(
sparse_tensor=frontier,
device=device,
)
candidate_batch_idx = new_batch_idx[has_space]
candidate_node_idx = new_node_idx[has_space]
candidate_batch_node_keys = new_batch_node_keys[has_space]
node_order = torch.argsort(candidate_node_idx, stable=True)
candidate_batch_idx = candidate_batch_idx[node_order]
candidate_node_idx = candidate_node_idx[node_order]
candidate_batch_node_keys = candidate_batch_node_keys[node_order]
batch_order = torch.argsort(candidate_batch_idx, stable=True)
sorted_batch_idx = candidate_batch_idx[batch_order]
sorted_node_idx = candidate_node_idx[batch_order]
sorted_batch_node_keys = candidate_batch_node_keys[batch_order]
n = sorted_batch_idx.size(0)
is_group_start = torch.zeros(n, dtype=torch.long, device=device)
is_group_start[0] = 1
if n > 1:
is_group_start[1:] = (sorted_batch_idx[1:] != sorted_batch_idx[:-1]).long()
group_id = is_group_start.cumsum(0) - 1
group_starts = torch.nonzero(is_group_start, as_tuple=True)[0]
positions = (
torch.arange(n, device=device)
- group_starts[group_id]
+ next_positions[sorted_batch_idx]
)
valid = positions < node_index_sequences.size(1)
if not valid.any().item():
return visited_batch_node_keys, _empty_sparse_tensor_like(
sparse_tensor=frontier,
device=device,
)
valid_batch_idx = sorted_batch_idx[valid]
valid_positions = positions[valid]
valid_node_idx = sorted_node_idx[valid]
node_index_sequences[valid_batch_idx, valid_positions] = valid_node_idx
valid_mask[valid_batch_idx, valid_positions] = True
inserted_counts = torch.zeros_like(next_positions)
inserted_counts.scatter_add_(
0,
valid_batch_idx,
torch.ones_like(valid_batch_idx, dtype=torch.long),
)
next_positions += inserted_counts
valid_batch_node_keys = sorted_batch_node_keys[valid]
valid_frontier = torch.sparse_coo_tensor(
torch.stack([valid_batch_idx, valid_node_idx]),
torch.ones(valid_node_idx.size(0), dtype=torch.float, device=device),
size=frontier.shape,
is_coalesced=True,
)
return (
torch.cat([visited_batch_node_keys, valid_batch_node_keys]).sort()[0],
valid_frontier,
)
def _build_sequence_layout_from_ppr_edges(
edge_index: Tensor,
edge_attr: Tensor,
anchor_indices: Tensor,
max_seq_len: int,
include_anchor_first: bool,
num_nodes: int,
device: torch.device,
return_edge_weights: bool = False,
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.
"""
batch_size = anchor_indices.size(0)
node_index_sequences = torch.full(
(batch_size, max_seq_len),
fill_value=-1,
dtype=torch.long,
device=device,
)
valid_mask = torch.zeros(
(batch_size, max_seq_len),
dtype=torch.bool,
device=device,
)
ppr_weight_sequences = None
if return_edge_weights:
ppr_weight_sequences = torch.zeros(
(batch_size, max_seq_len, 1),
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
valid_mask[:, 0] = True
start_pos = 1
else:
start_pos = 0
if start_pos >= max_seq_len:
return (
node_index_sequences,
valid_mask,
ppr_weight_sequences,
ppr_relation_feature_sequences,
)
edge_features = edge_attr.float()
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(
f"Requested '{PPR_RELATION_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),
dtype=torch.float,
device=device,
)
edge_weights = edge_features[:, 0]
relation_features = edge_features[:, 1:] if return_relation_features else None
anchor_batch_index_by_homo_idx = torch.full(
(num_nodes,),
fill_value=-1,
dtype=torch.long,
device=device,
)
anchor_batch_index_by_homo_idx[anchor_indices] = torch.arange(
batch_size, device=device
)
src_idx = edge_index[0]
dst_idx = edge_index[1]
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,
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,
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. Relation features, when
# present, are reordered by the same permutation to stay aligned with edges.
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]
if all_relation_features is not None:
all_relation_features = all_relation_features[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)
is_group_start[0] = 1
if n > 1:
is_group_start[1:] = (sorted_batch_idx[1:] != sorted_batch_idx[:-1]).long()
group_id = is_group_start.cumsum(0) - 1
group_starts = torch.nonzero(is_group_start, as_tuple=True)[0]
positions = torch.arange(n, device=device) - group_starts[group_id] + start_pos
valid = positions < max_seq_len
valid_batch_idx = sorted_batch_idx[valid]
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
if ppr_weight_sequences is not None:
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,
ppr_relation_feature_sequences,
)
def _gather_sequences_from_node_indices(
node_index_sequences: Tensor,
node_features: Tensor,
valid_mask: Tensor,
padding_value: float,
) -> Tensor:
"""Gather node features into padded sequences using precomputed node indices.
Args:
node_index_sequences: (batch_size, max_seq_len) node indices
node_features: (num_nodes, feature_dim) node features
valid_mask: (batch_size, max_seq_len) bool tensor indicating valid positions
padding_value: Value to use for padding
Returns:
(batch_size, max_seq_len, feature_dim) padded sequences.
"""
batch_size, max_seq_len = node_index_sequences.shape
feature_dim = node_features.size(-1)
if feature_dim == 0 or node_features.size(0) == 0:
return torch.full(
(batch_size, max_seq_len, feature_dim),
padding_value,
dtype=node_features.dtype,
device=node_features.device,
)
# A masked gather instead of bool-mask indexing: no nonzero, no host sync.
# Padding positions read rows spread over the whole table rather than one
# row, so the gather's backward does not pile every padding gradient (all
# zero) onto a single row.
spread_indices = (
torch.arange(batch_size * max_seq_len, device=node_index_sequences.device)
.view(batch_size, max_seq_len)
.remainder(node_features.size(0))
)
gathered = node_features[
torch.where(valid_mask, node_index_sequences, spread_indices)
]
return gathered.masked_fill(~valid_mask.unsqueeze(-1), padding_value)
def _lookup_anchor_relative_features(
anchor_indices: Tensor,
node_index_sequences: Tensor,
valid_mask: Tensor,
csr_matrices: Optional[list[Tensor]],
device: torch.device,
) -> Optional[Tensor]:
"""
Look up anchor-relative sparse values for each valid token in the sequence.
For each node in the sequence, this looks up the value PE[anchor_idx, node_idx]
from each provided sparse CSR matrix. This captures the relationship between
each sequence token and its anchor node (e.g., hop distance from anchor).
Args:
anchor_indices: (batch_size,) anchor node indices in homogeneous graph
node_index_sequences: (batch_size, max_seq_len) node indices for each sequence position
valid_mask: (batch_size, max_seq_len) bool tensor indicating valid positions
csr_matrices: List of sparse CSR matrices, each (num_nodes, num_nodes)
device: Device for output tensor
Returns:
features: (batch_size, max_seq_len, num_attrs) tensor where
features[b, i, k] = csr_matrices[k][anchor_indices[b], node_index_sequences[b, i]]
for valid positions, 0.0 for padding positions.
Returns None if csr_matrices is empty or None.
Example:
# batch_size=2, max_seq_len=4, num_attrs=1 (e.g., hop_distance)
# anchor_indices = [10, 20] (anchor nodes)
# node_index_sequences = [[10, 5, 7, -1], # anchor 10's sequence
# [20, 3, 8, 9]] # anchor 20's sequence
# valid_mask = [[T, T, T, F], [T, T, T, T]]
#
# Output shape: (2, 4, 1)
# features[0, :, 0] = [hop_dist[10,10], hop_dist[10,5], hop_dist[10,7], 0.0]
# features[1, :, 0] = [hop_dist[20,20], hop_dist[20,3], hop_dist[20,8], hop_dist[20,9]]
"""
if not csr_matrices:
return None
batch_size, max_seq_len = node_index_sequences.shape
num_attrs = len(csr_matrices)
features = torch.zeros(
(batch_size, max_seq_len, num_attrs),
dtype=torch.float,
device=device,
)
if not valid_mask.any():
return features
valid_batch_idx, valid_pos_idx = torch.nonzero(valid_mask, as_tuple=True)
valid_node_idx = node_index_sequences[valid_mask]
anchor_for_entry = anchor_indices[valid_batch_idx]
for attr_idx, pe_matrix in enumerate(csr_matrices):
pe_values = _lookup_csr_values(
csr_matrix=pe_matrix,
row_indices=anchor_for_entry,
col_indices=valid_node_idx,
)
features[valid_batch_idx, valid_pos_idx, attr_idx] = pe_values
return features
def _lookup_pairwise_relative_features(
node_index_sequences: Tensor,
valid_mask: Tensor,
csr_matrices: Optional[list[Tensor]],
attr_names: Optional[list[str]],
device: torch.device,
) -> tuple[Optional[Tensor], Optional[Tensor]]:
"""
Look up pairwise sparse values for each valid token pair in the sequence.
For each pair of nodes (i, j) in the sequence, this looks up the value
PE[node_i, node_j] from each provided sparse CSR matrix. This captures
pairwise relationships between all sequence tokens (e.g., random walk
structural encoding between any two nodes).
The output is typically used as attention bias in Graph Transformers,
added to attention scores before softmax.
Args:
node_index_sequences: (batch_size, max_seq_len) node indices for each sequence position
valid_mask: (batch_size, max_seq_len) bool tensor indicating valid positions
csr_matrices: List of sparse CSR matrices, each (num_nodes, num_nodes)
attr_names: Optional names for the pairwise attributes. Used only to
produce clearer error messages when multiple attrs disagree on
sparse support.
device: Device for output tensor
Returns:
features: (batch_size, max_seq_len, max_seq_len, num_attrs) tensor where
features[b, i, j, k] = csr_matrices[k][node_index_sequences[b, i], node_index_sequences[b, j]]
for valid (i, j) pairs, 0.0 for padding positions.
nonmissing_indices: (num_nonmissing_pairs, 3) long tensor containing
``(batch_idx, row_pos, col_pos)`` coordinates for valid diagonal
self pairs and valid sparse entries. Missing non-self pairs and
padding are omitted.
Returns (None, None) if csr_matrices is empty.
Example:
# batch_size=2, max_seq_len=3, num_attrs=1 (e.g., random_walk_se)
# node_index_sequences = [[10, 5, 7], # anchor 0's sequence
# [20, 3, -1]] # anchor 1's sequence (padded)
# valid_mask = [[T, T, T], [T, T, F]]
#
# Output shape: (2, 3, 3, 1)
#
# For batch 0 (all valid), features[0, :, :, 0] is a 3x3 matrix:
# node 10 node 5 node 7
# node 10 [PE[10,10], PE[10,5], PE[10,7]]
# node 5 [PE[5,10], PE[5,5], PE[5,7]]
# node 7 [PE[7,10], PE[7,5], PE[7,7]]
#
# For batch 1 (position 2 is padding), features[1, :, :, 0]:
# node 20 node 3 (pad)
# node 20 [PE[20,20], PE[20,3], 0.0]
# node 3 [PE[3,20], PE[3,3], 0.0]
# (pad) [0.0, 0.0, 0.0]
"""
if not csr_matrices:
return None, None
batch_size, max_seq_len = node_index_sequences.shape
num_attrs = len(csr_matrices)
features = torch.zeros(
(batch_size, max_seq_len, max_seq_len, num_attrs),
dtype=torch.float,
device=device,
)
pair_valid_mask = valid_mask.unsqueeze(2) & valid_mask.unsqueeze(1)
if not pair_valid_mask.any():
return features, torch.zeros((0, 3), dtype=torch.long, device=device)
valid_batch_indices, valid_row_positions, valid_col_positions = torch.nonzero(
pair_valid_mask,
as_tuple=True,
)
valid_row_indices = node_index_sequences[
valid_batch_indices,
valid_row_positions,
]
valid_col_indices = node_index_sequences[
valid_batch_indices,
valid_col_positions,
]
self_pair_mask = valid_row_positions == valid_col_positions
first_attr_name = attr_names[0] if attr_names else "attr_0"
nonmissing_support: Optional[Tensor] = None
for attr_idx, pe_matrix in enumerate(csr_matrices):
pe_values, found_mask = _lookup_csr_values_and_found(
csr_matrix=pe_matrix,
row_indices=valid_row_indices,
col_indices=valid_col_indices,
)
features[
valid_batch_indices,
valid_row_positions,
valid_col_positions,
attr_idx,
] = pe_values
attr_nonmissing_support = found_mask | self_pair_mask
if attr_idx == 0:
nonmissing_support = attr_nonmissing_support
continue
if nonmissing_support is None or not torch.equal(
nonmissing_support,
attr_nonmissing_support,
):
attr_name = attr_names[attr_idx] if attr_names else f"attr_{attr_idx}"
raise ValueError(
"Pairwise attention bias attributes must share identical "
"nonmissing support after treating valid diagonal self pairs "
f"as nonmissing, but '{first_attr_name}' and '{attr_name}' "
"differ."
)
assert nonmissing_support is not None
pairwise_nonmissing_indices = torch.stack(
[
valid_batch_indices[nonmissing_support],
valid_row_positions[nonmissing_support],
valid_col_positions[nonmissing_support],
],
dim=1,
)
return features, pairwise_nonmissing_indices
def _lookup_pairwise_relation_indices(
edge_index_dict: dict[EdgeType, Tensor],
node_index_sequences: Tensor,
valid_mask: Tensor,
relation_edge_types: Optional[list[GiGLEdgeType]],
node_type_offsets: dict[NodeType, int],
num_nodes: int,
device: torch.device,
expanded_token_mask: Optional[Tensor] = None,
relation_channel_masks: Optional[dict[EdgeType, int]] = None,
frontier_batch_node_keys: Optional[Tensor] = None,
sampling_direction: Literal["in", "out"] = "out",
) -> Optional[Tensor]:
"""Build sparse relation coordinates for valid token pairs.
For a directed edge ``source -> target``, attention uses ``query=target`` and
``key=source`` so relation-aware attention follows message-passing
orientation.
If ``frontier_batch_node_keys`` is given (sorted ``batch_idx * num_nodes +
node_idx`` keys), an edge is kept only when its expanded endpoint is in
that anchor's frontier. K-hop sequences pass this so relation edges match
the anchor's own sample; PPR sequences pass ``None``.
If ``expanded_token_mask`` is given (a per-token bitmask of the PPR channels
that pushed each token's node), an edge is kept only when the anchor expanded
its owner -- the query/target position, since GLT reverses sampled edges in
collation. For typed PPR, ``relation_channel_masks`` maps an edge type to the
channels allowed to fetch it, so the edge is kept only when the owner was
pushed in one of those channels. PPR sequences pass these; k-hop passes
``None``.
"""
if not relation_edge_types:
return None
token_occurrences = _build_token_occurrence_index(
node_index_sequences=node_index_sequences,
valid_mask=valid_mask,
num_nodes=num_nodes,
device=device,
)
if token_occurrences.batch_indices.numel() == 0:
return torch.zeros((0, 4), dtype=torch.long, device=device)
relation_index_parts: list[Tensor] = []
for relation_idx, edge_type in enumerate(relation_edge_types):
edge_type_tuple = edge_type.tuple_repr()
if edge_type_tuple not in edge_index_dict:
continue
edge_index = edge_index_dict[edge_type_tuple].to(
device=device,
dtype=torch.long,
)
if edge_index.numel() == 0:
continue
src_offset = int(node_type_offsets[edge_type.src_node_type])
dst_offset = int(node_type_offsets[edge_type.dst_node_type])
source_indices = edge_index[0] + src_offset
target_indices = edge_index[1] + dst_offset
(
relation_batch_indices,
relation_query_positions,
relation_key_positions,
) = _match_directed_edges_to_token_pairs(
source_indices=source_indices,
target_indices=target_indices,
token_occurrences=token_occurrences,
num_nodes=num_nodes,
device=device,
)
if frontier_batch_node_keys is not None:
expanded_positions = (
relation_query_positions
if sampling_direction == "in"
else relation_key_positions
)
expanded_batch_node_keys = (
relation_batch_indices * num_nodes
+ node_index_sequences[relation_batch_indices, expanded_positions]
)
in_frontier = _sorted_contains(
sorted_values=frontier_batch_node_keys,
values=expanded_batch_node_keys,
)
relation_batch_indices = relation_batch_indices[in_frontier]
relation_query_positions = relation_query_positions[in_frontier]
relation_key_positions = relation_key_positions[in_frontier]
if relation_batch_indices.numel() == 0:
continue
relation_channel_mask = (
relation_channel_masks.get(edge_type.tuple_repr())
if relation_channel_masks is not None
else None
)
if expanded_token_mask is not None:
# Keep a relation only when the anchor's own push expanded the OWNER --
# the fetched node, which ppr_is_expanded describes. GLT reverses sampled
# edges in collation, so the owner is the target (query) position, not
# the source/key. For typed PPR the mask is a per-node channel bitmask:
# keep the edge only if the anchor pushed the owner in a channel whose
# allowlist contains this edge type (relation_channel_mask). Untyped PPR
# has no channel mask, so any nonzero (pushed) owner passes.
owner_mask = expanded_token_mask[
relation_batch_indices, relation_query_positions
]
if relation_channel_mask is None:
keep = owner_mask != 0
else:
keep = (owner_mask & relation_channel_mask) != 0
relation_batch_indices = relation_batch_indices[keep]
relation_query_positions = relation_query_positions[keep]
relation_key_positions = relation_key_positions[keep]
if relation_batch_indices.numel() == 0:
continue
relation_indices = torch.stack(
[
relation_batch_indices,
relation_query_positions,
relation_key_positions,
torch.full(
(relation_batch_indices.size(0),),
relation_idx,
dtype=torch.long,
device=device,
),
],
dim=1,
)
relation_index_parts.append(torch.unique(relation_indices, dim=0))
if not relation_index_parts:
return torch.zeros((0, 4), dtype=torch.long, device=device)
return torch.cat(relation_index_parts, dim=0)
def _get_sampled_frontier_batch_node_keys(
anchor_indices: Tensor,
edge_index: Tensor,
num_nodes: int,
hop_distance: int,
device: torch.device,
) -> Tensor:
"""Return sorted ``batch_idx * num_nodes + node_idx`` keys for nodes fewer
than ``hop_distance`` hops from each anchor.
These are the nodes a k-hop sampler expands for that anchor, so their
sampled edges match the anchor's own sample. ``edge_index`` must already
point in the traversal direction.
"""
if hop_distance <= 0:
return torch.zeros((0,), dtype=torch.long, device=device)
frontier = _get_k_hop_neighbors_sparse(
anchor_indices=anchor_indices,
edge_index=edge_index,
num_nodes=num_nodes,
k=hop_distance - 1,
device=device,
)
batch_idx, node_idx = frontier.indices()
# Coalesced COO indices are row-major sorted, so the keys are sorted too.
return batch_idx * num_nodes + node_idx
def _sorted_contains(sorted_values: Tensor, values: Tensor) -> Tensor:
"""Return a bool mask of which ``values`` appear in ``sorted_values``."""
if sorted_values.numel() == 0:
return torch.zeros_like(values, dtype=torch.bool)
positions = torch.searchsorted(sorted_values, values).clamp(
max=sorted_values.numel() - 1
)
return sorted_values[positions] == values
def _build_token_occurrence_index(
node_index_sequences: Tensor,
valid_mask: Tensor,
num_nodes: int,
device: torch.device,
) -> _TokenOccurrenceIndex:
"""Index valid sequence tokens for sparse directed-edge to token matching."""
token_batch_indices, token_positions = torch.nonzero(valid_mask, as_tuple=True)
token_batch_indices = token_batch_indices.to(device=device, dtype=torch.long)
token_positions = token_positions.to(device=device, dtype=torch.long)
token_node_indices = node_index_sequences[token_batch_indices, token_positions].to(
device=device,
dtype=torch.long,
)
sorted_token_node_indices, node_sort_perm = torch.sort(token_node_indices)
token_batch_node_keys = token_batch_indices * num_nodes + token_node_indices
sorted_token_batch_node_keys, batch_node_sort_perm = torch.sort(
token_batch_node_keys
)
return _TokenOccurrenceIndex(
batch_indices=token_batch_indices,
positions=token_positions,
node_indices=token_node_indices,
sorted_node_indices=sorted_token_node_indices,
node_sort_perm=node_sort_perm,
sorted_batch_node_keys=sorted_token_batch_node_keys,
batch_node_sort_perm=batch_node_sort_perm,
)
def _match_directed_edges_to_token_pairs(
source_indices: Tensor,
target_indices: Tensor,
token_occurrences: _TokenOccurrenceIndex,
num_nodes: int,
device: torch.device,
) -> tuple[Tensor, Tensor, Tensor]:
"""Map ``source -> target`` graph edges onto valid sequence coordinates."""
empty = torch.zeros((0,), dtype=torch.long, device=device)
if source_indices.numel() == 0 or token_occurrences.batch_indices.numel() == 0:
return empty, empty, empty
source_indices = source_indices.to(device=device, dtype=torch.long)
target_indices = target_indices.to(device=device, dtype=torch.long)
target_lower_bounds = torch.searchsorted(
token_occurrences.sorted_node_indices,
target_indices,
right=False,
)
target_upper_bounds = torch.searchsorted(
token_occurrences.sorted_node_indices,
target_indices,
right=True,
)
target_match_counts = target_upper_bounds - target_lower_bounds
matched_edge_mask = target_match_counts > 0
if not matched_edge_mask.any():
return empty, empty, empty
matched_edge_indices = torch.nonzero(matched_edge_mask, as_tuple=True)[0]
matched_target_counts = target_match_counts[matched_edge_indices]
total_target_matches = int(matched_target_counts.sum().item())
repeated_target_edge_indices = torch.repeat_interleave(
matched_edge_indices,
matched_target_counts,
)
repeated_target_lower_bounds = torch.repeat_interleave(
target_lower_bounds[matched_edge_indices],
matched_target_counts,
)
target_group_start_offsets = torch.repeat_interleave(
torch.cumsum(matched_target_counts, dim=0) - matched_target_counts,
matched_target_counts,
)
target_sorted_positions = (
repeated_target_lower_bounds
+ torch.arange(total_target_matches, device=device, dtype=torch.long)
- target_group_start_offsets
)
target_token_indices = token_occurrences.node_sort_perm[target_sorted_positions]
target_batch_indices = token_occurrences.batch_indices[target_token_indices]
target_query_positions = token_occurrences.positions[target_token_indices]
source_query_keys = (
target_batch_indices * num_nodes + source_indices[repeated_target_edge_indices]
)
source_lower_bounds = torch.searchsorted(
token_occurrences.sorted_batch_node_keys,
source_query_keys,
right=False,
)
source_upper_bounds = torch.searchsorted(
token_occurrences.sorted_batch_node_keys,
source_query_keys,
right=True,
)
source_match_counts = source_upper_bounds - source_lower_bounds
matched_target_mask = source_match_counts > 0
if not matched_target_mask.any():
return empty, empty, empty
matched_target_indices = torch.nonzero(matched_target_mask, as_tuple=True)[0]
matched_source_counts = source_match_counts[matched_target_indices]
total_source_matches = int(matched_source_counts.sum().item())
repeated_target_indices = torch.repeat_interleave(
matched_target_indices,
matched_source_counts,
)
repeated_source_lower_bounds = torch.repeat_interleave(
source_lower_bounds[matched_target_indices],
matched_source_counts,
)
source_group_start_offsets = torch.repeat_interleave(
torch.cumsum(matched_source_counts, dim=0) - matched_source_counts,
matched_source_counts,
)
source_sorted_positions = (
repeated_source_lower_bounds
+ torch.arange(total_source_matches, device=device, dtype=torch.long)
- source_group_start_offsets
)
source_token_indices = token_occurrences.batch_node_sort_perm[
source_sorted_positions
]
return (
target_batch_indices[repeated_target_indices],
target_query_positions[repeated_target_indices],
token_occurrences.positions[source_token_indices],
)
def _get_k_hop_neighbors_sparse(
anchor_indices: Tensor,
edge_index: Tensor,
num_nodes: int,
k: int,
device: torch.device,
) -> Tensor:
"""
Get k-hop reachable nodes for all anchors using sparse matrix multiplication.
Follows the same efficient pattern as AddHeteroRandomWalkPE: simple sparse
matrix powers to accumulate reachability without expensive membership checks.
Args:
anchor_indices: (batch_size,) anchor node indices
edge_index: (2, num_edges) edge index
num_nodes: Total number of nodes
k: Number of hops
device: Device for tensors
Returns:
reachable: (batch_size, num_nodes) sparse tensor, non-zero if node is reachable within k hops
"""
batch_size = anchor_indices.size(0)
# Build sparse adjacency matrix (binarized) - coalesce once
adj = _binarize_sparse_tensor(
sparse_tensor=to_torch_sparse_tensor(edge_index, size=(num_nodes, num_nodes)),
device=device,
)
# Initialize: sparse matrix where row i has a 1 at column anchor_indices[i]
reachable = torch.sparse_coo_tensor(
torch.stack(
[
torch.arange(batch_size, device=device),
anchor_indices,
]
),
torch.ones(batch_size, device=device, dtype=torch.float),
size=(batch_size, num_nodes),
is_coalesced=True,
)
current = reachable
for _ in range(k):
# Expand: current @ adj gives nodes reachable in one more hop
current = torch.sparse.mm(current, adj)
if current._nnz() == 0:
break
# Accumulate into reachable
reachable = reachable + current
# Coalesce and binarize final result
return _binarize_sparse_tensor(sparse_tensor=reachable, device=device)
def _lookup_csr_values(
csr_matrix: Tensor,
row_indices: Tensor,
col_indices: Tensor,
default_value: float = 0.0,
) -> Tensor:
"""
Look up values in a CSR sparse matrix for given (row, col) pairs.
Vectorized CSR lookup: for each query, slice the row and search for the column.
Time complexity: O(n * avg_nnz_per_row), typically O(n) for sparse matrices.
Args:
csr_matrix: (num_rows, num_cols) sparse CSR tensor
row_indices: (n,) row indices to look up
col_indices: (n,) column indices to look up
default_value: Value for missing entries (default: 0.0)
Returns:
(n,) values from csr_matrix[row, col], or default_value if not present
"""
values, _ = _lookup_csr_values_and_found(
csr_matrix=csr_matrix,
row_indices=row_indices,
col_indices=col_indices,
default_value=default_value,
)
return values
def _lookup_csr_values_and_found(
csr_matrix: Tensor,
row_indices: Tensor,
col_indices: Tensor,
default_value: float = 0.0,
) -> tuple[Tensor, Tensor]:
"""
Look up values in a CSR sparse matrix and report which entries were present.
Returns both the looked-up values and a boolean found-mask so callers can
distinguish missing sparse entries from explicit zero-valued entries.
"""
n = row_indices.size(0)
device = row_indices.device
if n == 0:
return (
torch.zeros(0, device=device, dtype=torch.float),
torch.zeros(0, device=device, dtype=torch.bool),
)
crow_indices = csr_matrix.crow_indices()
col_indices_csr = csr_matrix.col_indices()
values_csr = csr_matrix.values()
# Get row start/end pointers
row_starts = crow_indices[row_indices]
row_ends = crow_indices[row_indices + 1]
row_lengths = row_ends - row_starts
max_row_len = row_lengths.max().item()
if max_row_len == 0:
return (
torch.full((n,), default_value, device=device, dtype=torch.float),
torch.zeros((n,), device=device, dtype=torch.bool),
)
# Build offset matrix: (n, max_row_len)
offsets = row_starts.unsqueeze(1) + torch.arange(max_row_len, device=device)
valid_mask = offsets < row_ends.unsqueeze(1)
# Safe indexing with clamping
nnz = col_indices_csr.size(0)
offsets_clamped = offsets.clamp(max=max(nnz - 1, 0))
# Get columns at offsets and find matches
cols_at_offsets = col_indices_csr[offsets_clamped]
col_matches = (cols_at_offsets == col_indices.unsqueeze(1)) & valid_mask
# Find which queries have matches
found = col_matches.any(dim=1)
# Initialize output
result = torch.full((n,), default_value, device=device, dtype=torch.float)
if found.any():
# Get match positions and retrieve values
match_offsets = col_matches.float().argmax(dim=1)
value_indices = row_starts[found] + match_offsets[found]
result[found] = values_csr[value_indices].float()
return result, found
