"""Relation-aware inductive GraphSAGE and optional GAT encoders.""" from __future__ import annotations from dataclasses import dataclass import torch from torch import Tensor, nn from torch.nn import functional from graph.config import GraphConfig from graph.types import GraphTensors @dataclass(frozen=True) class GraphEncoderOutput: node_states: Tensor relation_states: Tensor edge_states: Tensor attention_weights: Tensor | None def _edge_query( relation_states: Tensor, query_state: Tensor | None, edge_query: Tensor | None ) -> Tensor: """The query vector each edge's message is conditioned on. ``edge_query`` gives one row per edge, which is how a batch of disjoint graphs, each with its own query, is encoded in a single pass. """ if edge_query is not None: return edge_query if query_state is None: return torch.zeros_like(relation_states) return query_state.unsqueeze(0).expand_as(relation_states) class RelationAwareSAGELayer(nn.Module): def __init__( self, hidden_dim: int, edge_dim: int, num_relations: int, dropout: float, residual: bool = False, ) -> None: super().__init__() self.residual = residual self.self_proj = nn.Linear(hidden_dim, hidden_dim) self.message_proj = nn.Linear( hidden_dim * 2 + hidden_dim + edge_dim, hidden_dim ) self.relation_embedding = nn.Embedding(num_relations, hidden_dim) self.norm = nn.LayerNorm(hidden_dim) self.dropout = nn.Dropout(dropout) def _skip(self, node_states: Tensor) -> Tensor | float: """The identity path, when enabled; it keeps depth from oversmoothing.""" return node_states if self.residual else 0.0 def forward( self, node_states: Tensor, graph: GraphTensors, relation_ids: Tensor, query_state: Tensor | None = None, edge_query: Tensor | None = None, ) -> tuple[Tensor, Tensor, Tensor | None]: source, target = graph.edge_index source_states = node_states[source] target_states = node_states[target] relation_states = self.relation_embedding(relation_ids) query_for_edges = _edge_query(relation_states, query_state, edge_query) messages = self.message_proj( torch.cat( (source_states, target_states, relation_states, graph.edge_features), dim=-1, ) ) messages = messages + query_for_edges aggregate = torch.zeros_like(node_states) aggregate.index_add_(0, target, messages) degree = torch.zeros( (graph.num_nodes,), device=node_states.device, dtype=node_states.dtype ) degree.index_add_(0, target, torch.ones_like(target, dtype=node_states.dtype)) aggregate = aggregate / degree.clamp_min(1.0).unsqueeze(-1) updated = self.norm( self._skip(node_states) + self.self_proj(node_states) + self.dropout(functional.gelu(aggregate)) ) return updated, relation_states, None class RelationAwareGATLayer(RelationAwareSAGELayer): """GAT ablation that retains per-edge attention for explanation.""" def __init__( self, hidden_dim: int, edge_dim: int, num_relations: int, dropout: float, residual: bool = False, ) -> None: super().__init__(hidden_dim, edge_dim, num_relations, dropout, residual) self.attn = nn.Linear(hidden_dim * 2 + hidden_dim + edge_dim, 1) def forward( self, node_states: Tensor, graph: GraphTensors, relation_ids: Tensor, query_state: Tensor | None = None, edge_query: Tensor | None = None, ) -> tuple[Tensor, Tensor, Tensor | None]: source, target = graph.edge_index source_states = node_states[source] target_states = node_states[target] relation_states = self.relation_embedding(relation_ids) query_for_edges = _edge_query(relation_states, query_state, edge_query) edge_input = torch.cat( (source_states, target_states, relation_states, graph.edge_features), dim=-1 ) logits = self.attn(edge_input).squeeze(-1) + ( relation_states * query_for_edges ).sum(-1) / (node_states.shape[-1] ** 0.5) weights = torch.zeros_like(logits) for node in range(graph.num_nodes): mask = target == node if mask.any(): weights[mask] = torch.softmax(logits[mask], dim=0) messages = self.message_proj(edge_input) * weights.unsqueeze(-1) aggregate = torch.zeros_like(node_states) aggregate.index_add_(0, target, messages) updated = self.norm( self._skip(node_states) + self.self_proj(node_states) + self.dropout(functional.gelu(aggregate)) ) return updated, relation_states, weights class GraphEncoder(nn.Module): """Inductive graph encoder; no trainable closed-world node-ID table is used.""" def __init__(self, config: GraphConfig) -> None: super().__init__() self.config = config self.node_type_embedding = nn.Embedding( config.num_node_types, config.hidden_dim ) self.node_input = nn.Linear(config.graph_input_dim, config.hidden_dim) layer_cls = ( RelationAwareGATLayer if config.graph_encoder == "gat" or config.use_gat else RelationAwareSAGELayer ) self.layers = nn.ModuleList( [ layer_cls( config.hidden_dim, config.edge_input_dim, config.num_relations, config.dropout, residual=config.graph_residual, ) for _ in range(config.graph_layers) ] ) self.edge_input = nn.Linear(config.edge_input_dim, config.hidden_dim) self.final_norm = nn.LayerNorm(config.hidden_dim) self.start_embedding = ( nn.Embedding(2, config.hidden_dim) if config.start_marking else None ) def forward( self, graph: GraphTensors, query_state: Tensor | None = None, *, edge_query: Tensor | None = None, start_mask: Tensor | None = None, node_query: Tensor | None = None, ) -> GraphEncoderOutput: """Encode ``graph``, conditioned on the query where it is given. ``edge_query``/``node_query`` give one query row per edge/node, for a batch of graphs with different queries; otherwise ``query_state`` is shared by the whole graph. ``start_mask`` flags the query's start entities and is used when the encoder was built with ``start_marking``. """ # External graph services are open-world: a tenant may introduce more # node/relation labels than the configured embedding tables, so the stable # global identities are reduced into the learned vocabulary here. The # graph itself is never rewritten: its *local* type indices are what the # caller uses to pool per-edge states into a per-graph relation bank and to # map a bank row back to a relation name, and overwriting them with reduced # hashes would corrupt that bookkeeping. node_type_ids = graph.node_type_identity.remainder(self.config.num_node_types) edge_type_ids = graph.edge_type_identity.remainder(self.config.num_relations) node_states = self.node_input(graph.node_features) + self.node_type_embedding( node_type_ids ) if self.start_embedding is not None and start_mask is not None: marked = start_mask.to(torch.long) node_states = node_states + self.start_embedding(marked) if node_query is None and query_state is not None: node_query = query_state.reshape(1, -1).expand_as(node_states) if node_query is not None: node_states = node_states + marked.unsqueeze(-1).to( node_states.dtype ) * node_query.to(node_states.dtype) all_relation_states: list[Tensor] = [] attention: Tensor | None = None for layer in self.layers: node_states, relation_states, layer_attention = layer( node_states, graph, edge_type_ids, query_state, edge_query ) if layer_attention is not None: attention = layer_attention all_relation_states.append(relation_states) node_states = self.final_norm(node_states) edge_states = self.edge_input(graph.edge_features) relation_states = ( all_relation_states[-1] if all_relation_states else torch.empty((0, self.config.hidden_dim), device=node_states.device) ) return GraphEncoderOutput(node_states, relation_states, edge_states, attention)