"""PAT: query-conditioned, relation-aware Path-Aware Transformer.""" from __future__ import annotations import itertools from collections.abc import Mapping, Sequence from dataclasses import dataclass from enum import IntEnum import torch from torch import Tensor, nn from torch.nn import functional from torch.utils.checkpoint import checkpoint from graph.config import PATConfig from graph.types import ConstraintState, GraphPath from graph.utils import qk_normalize, soft_cap, stable_hash class PATTokenType(IntEnum): QUERY = 0 ENTITY = 1 RELATION = 2 EDGE = 3 TYPE = 4 PATH_START = 5 PATH_STEP = 6 PATH_END = 7 HOP = 8 CONSTRAINT = 9 @dataclass(frozen=True) class PATInput: token_features: Tensor token_type_ids: Tensor hop_positions: Tensor attention_mask: Tensor @dataclass(frozen=True) class PATOutput: path_states: Tensor relation_sequence_states: Tensor transition_states: Tensor compatibility: Tensor token_states: Tensor endpoint_states: Tensor hop_states: Tensor constraint_compatibility: Tensor class RelationAwareAttention(nn.Module): def __init__( self, hidden_dim: int, num_heads: int, num_token_types: int, dropout: float, logit_cap: float = 0.0, qk_norm: bool = False, ) -> None: super().__init__() if hidden_dim % num_heads: raise ValueError("PAT hidden_dim must be divisible by num_heads") self.logit_cap = logit_cap self.qk_norm = qk_norm self.hidden_dim = hidden_dim self.num_heads = num_heads self.head_dim = hidden_dim // num_heads self.q_proj = nn.Linear(hidden_dim, hidden_dim) self.k_proj = nn.Linear(hidden_dim, hidden_dim) self.v_proj = nn.Linear(hidden_dim, hidden_dim) self.out_proj = nn.Linear(hidden_dim, hidden_dim) self.pair_bias = nn.Parameter( torch.zeros(num_token_types, num_token_types, num_heads) ) self.dropout = nn.Dropout(dropout) def forward( self, states: Tensor, token_type_ids: Tensor, mask: Tensor, query_state: Tensor ) -> Tensor: batch, length, _ = states.shape q = ( self.q_proj(states) .view(batch, length, self.num_heads, self.head_dim) .transpose(1, 2) ) k = ( self.k_proj(states) .view(batch, length, self.num_heads, self.head_dim) .transpose(1, 2) ) v = ( self.v_proj(states) .view(batch, length, self.num_heads, self.head_dim) .transpose(1, 2) ) if self.qk_norm: q, k = qk_normalize(q), qk_normalize(k) logits = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim**0.5) pair = self.pair_bias[ token_type_ids.unsqueeze(-1), token_type_ids.unsqueeze(-2) ] logits = logits + pair.permute(0, 3, 1, 2) query_key = query_state.view(batch, self.num_heads, self.head_dim) logits = logits + torch.einsum("bhd,bhld->bhl", query_key, k).unsqueeze(-2) / ( self.head_dim**0.5 ) logits = soft_cap(logits, self.logit_cap) valid = mask.to(torch.bool) logits = logits.masked_fill( ~valid[:, None, None, :], torch.finfo(logits.dtype).min ) weights = torch.softmax(logits, dim=-1).to(v.dtype) weights = self.dropout(weights) result = ( torch.matmul(weights, v) .transpose(1, 2) .contiguous() .view(batch, length, self.hidden_dim) ) return self.out_proj(result) class PATBlock(nn.Module): def __init__( self, config: PATConfig, logit_cap: float = 0.0, qk_norm: bool = False ) -> None: super().__init__() self.attn = RelationAwareAttention( config.hidden_dim, config.num_heads, config.num_token_types, config.dropout, logit_cap, qk_norm, ) self.norm1 = nn.LayerNorm(config.hidden_dim) self.norm2 = nn.LayerNorm(config.hidden_dim) self.ff = nn.Sequential( nn.Linear(config.hidden_dim, config.hidden_dim * 4), nn.GELU(), nn.Dropout(config.dropout), nn.Linear(config.hidden_dim * 4, config.hidden_dim), ) self.dropout = nn.Dropout(config.dropout) def forward( self, states: Tensor, token_type_ids: Tensor, mask: Tensor, query_state: Tensor ) -> Tensor: attended = self.attn(self.norm1(states), token_type_ids, mask, query_state) states = states + self.dropout(attended) return states + self.dropout(self.ff(self.norm2(states))) class PathAwareTransformer(nn.Module): """PAT operates on structured vectors, never on a textual graph dump.""" def __init__( self, config: PATConfig, gradient_checkpointing: bool = False, logit_cap: float = 0.0, qk_norm: bool = False, ) -> None: super().__init__() self.config = config self.gradient_checkpointing = gradient_checkpointing self.token_type_embedding = nn.Embedding( config.num_token_types, config.hidden_dim ) self.path_position_embedding = nn.Embedding( config.max_path_length * 5 + 5, config.hidden_dim ) self.hop_embedding = nn.Embedding(config.max_path_length + 2, config.hidden_dim) self.type_embedding = nn.Embedding(config.type_vocab_size, config.hidden_dim) self.attribute_embedding = nn.Embedding( config.attribute_vocab_size, config.hidden_dim ) self.query_gate = nn.Linear(config.hidden_dim, config.hidden_dim) self.blocks = nn.ModuleList( [PATBlock(config, logit_cap, qk_norm) for _ in range(config.num_layers)] ) self.final_norm = nn.LayerNorm(config.hidden_dim) self.query_projection = nn.Linear(config.hidden_dim, config.hidden_dim) def forward(self, inputs: PATInput, query_state: Tensor) -> PATOutput: token_types = inputs.token_type_ids.clamp(0, self.config.num_token_types - 1) positions = torch.arange( inputs.token_features.shape[1], device=inputs.token_features.device ) positions = positions.clamp_max(self.config.max_path_length * 5 + 4).unsqueeze( 0 ) hops = inputs.hop_positions.clamp(0, self.config.max_path_length + 1) states = inputs.token_features + self.token_type_embedding(token_types) if query_state.ndim == 1: query_state = query_state.unsqueeze(0) if query_state.shape[0] == 1 and states.shape[0] != 1: # contiguous(): MPS aborts on Linear over a zero-stride expanded input query_state = query_state.expand(states.shape[0], -1).contiguous() states = ( states + self.path_position_embedding(positions) + self.hop_embedding(hops) ) attention_query = ( query_state if self.config.query_conditioning else torch.zeros_like(query_state) ) if self.config.query_conditioning: states = states + self.query_gate(query_state).unsqueeze(1) for block in self.blocks: if self.gradient_checkpointing and self.training: states = checkpoint( lambda values, layer=block: layer( values, token_types, inputs.attention_mask, attention_query ), states, use_reentrant=False, ) else: states = block( states, token_types, inputs.attention_mask, attention_query ) states = self.final_norm(states) mask = inputs.attention_mask.to(states.dtype).unsqueeze(-1) path_states = (states * mask).sum(1) / mask.sum(1).clamp_min(1.0) path_states = functional.normalize(path_states, dim=-1) relation_mask = ( token_types == int(PATTokenType.RELATION) ) & inputs.attention_mask.bool() max_relations = max(1, int(relation_mask.sum(dim=1).max().item())) # Each row's relation tokens, left-aligned and zero-padded: one scatter # instead of a gather-and-pad per path. rows, columns = relation_mask.nonzero(as_tuple=True) slots = relation_mask.long().cumsum(1)[rows, columns] - 1 relation_sequence_states = states.new_zeros( (states.shape[0], max_relations, states.shape[-1]) ).index_put((rows, slots), states[rows, columns]) transition_states = relation_sequence_states compatibility = ( path_states * functional.normalize(self.query_projection(attention_query), dim=-1) ).sum(-1) entity_mask = ( token_types == int(PATTokenType.ENTITY) ) & inputs.attention_mask.bool() entity_positions = ( torch.arange(states.shape[1], device=states.device) .unsqueeze(0) .expand(states.shape[0], -1) ) last_entity = ( torch.where( entity_mask, entity_positions, torch.zeros_like(entity_positions) ) .max(dim=1) .values ) batch_indices = torch.arange(states.shape[0], device=states.device) endpoint_states = states[batch_indices, last_entity] hop_states = relation_sequence_states.mean(1) constraint_compatibility = functional.cosine_similarity( path_states, self.query_projection(attention_query), dim=-1 ) return PATOutput( path_states, relation_sequence_states, transition_states, compatibility, states, endpoint_states, hop_states, constraint_compatibility, ) def build_input( self, paths: Sequence[GraphPath], node_states: Mapping[str, Tensor], relation_states: Mapping[str, Tensor], query_state: Tensor, constraint_state: ConstraintState | None = None, ) -> PATInput: if query_state.ndim == 2: query_state = query_state[0] if not paths: raise ValueError("PAT requires at least one candidate path") feature_rows: list[Tensor] = [] type_rows: list[Tensor] = [] hop_rows: list[Tensor] = [] for path in paths: if not path.steps: raise ValueError("PAT cannot encode an empty path") features = [torch.zeros_like(query_state)] types = [int(PATTokenType.PATH_START)] hops = [0] if constraint_state is not None: constraint_signal = float( constraint_state.max_hops - constraint_state.min_hops ) if constraint_state.exact_hops is not None: constraint_signal += float(constraint_state.exact_hops) features.append( query_state * (1.0 + constraint_signal / max(1.0, self.config.max_path_length)) ) types.append(int(PATTokenType.CONSTRAINT)) hops.append(0) first = path.steps[0] features.append(node_states[first.source_id]) types.append(int(PATTokenType.ENTITY)) hops.append(0) features.append(self._type_state(first.source_type, query_state)) types.append(int(PATTokenType.TYPE)) hops.append(0) for hop, step in enumerate(path.steps, start=1): features.append( relation_states.get(step.relation_id, torch.zeros_like(query_state)) ) types.append(int(PATTokenType.RELATION)) hops.append(hop) features.append( self._attribute_state(step.edge_attributes, query_state) ) types.append(int(PATTokenType.EDGE)) hops.append(hop) features.append(node_states[step.target_id]) types.append(int(PATTokenType.ENTITY)) hops.append(hop) features.append(self._type_state(step.target_type, query_state)) types.append(int(PATTokenType.TYPE)) hops.append(hop) features.append( query_state * (float(hop) / max(1.0, self.config.max_path_length)) ) types.append(int(PATTokenType.HOP)) hops.append(hop) features.append(torch.zeros_like(query_state)) types.append(int(PATTokenType.PATH_END)) hops.append(len(path.steps)) feature_rows.append(torch.stack(features)) type_rows.append( torch.tensor(types, dtype=torch.long, device=query_state.device) ) hop_rows.append( torch.tensor(hops, dtype=torch.long, device=query_state.device) ) length = max(row.shape[0] for row in feature_rows) hidden = query_state.shape[-1] padded_features = torch.zeros( (len(paths), length, hidden), dtype=query_state.dtype, device=query_state.device, ) padded_types = torch.full( (len(paths), length), int(PATTokenType.PATH_END), dtype=torch.long, device=query_state.device, ) padded_hops = torch.zeros( (len(paths), length), dtype=torch.long, device=query_state.device ) padded_mask = torch.zeros( (len(paths), length), dtype=torch.long, device=query_state.device ) for i, (feature_row, type_row, hop_row) in enumerate( zip(feature_rows, type_rows, hop_rows, strict=True) ): size = feature_row.shape[0] padded_features[i, :size] = feature_row padded_types[i, :size] = type_row padded_hops[i, :size] = hop_row padded_mask[i, :size] = 1 return PATInput(padded_features, padded_types, padded_hops, padded_mask) def build_batch_input( self, paths: Sequence[GraphPath], owners: Sequence[int], node_rows: Sequence[Mapping[str, int]], relation_rows: Sequence[Mapping[str, int]], node_table: Tensor, relation_table: Tensor, query_states: Tensor, constraints: Sequence[ConstraintState | None], ) -> PATInput: """:meth:`build_input` for paths drawn from many examples at once. Path ``i`` belongs to example ``owners[i]``, whose node and relation states are rows ``node_rows[owner][id]`` of ``node_table`` and ``relation_rows[owner][id]`` of ``relation_table``. Every token is a row of one feature table plus a multiple of its example's query state, so the whole batch is assembled by a single gather instead of one small tensor operation per token. Token for token the result matches :meth:`build_input`, padded to the batch's longest path. """ if not paths: raise ValueError("PAT requires at least one candidate path") max_length = max(1.0, self.config.max_path_length) type_ids: dict[int, int] = {} attribute_sets: dict[tuple[int, ...], int] = {} node_offset = 1 relation_offset = node_offset + node_table.shape[0] # Tokens reference types and attribute sets by position until their rows # are known; ("t", k) and ("a", k) are resolved after the table is built. token_rows: list[list[int | tuple[str, int]]] = [] token_scales: list[list[float]] = [] token_types: list[list[int]] = [] token_hops: list[list[int]] = [] def type_slot(type_name: str | None) -> tuple[str, int]: index = stable_hash(type_name or "unknown", self.config.type_vocab_size) return ("t", type_ids.setdefault(index, len(type_ids))) def attribute_slot(attributes: Mapping[str, str]) -> int | tuple[str, int]: if not attributes: return 0 key = tuple( stable_hash(f"{name}={value}", self.config.attribute_vocab_size) for name, value in sorted(attributes.items()) ) return ("a", attribute_sets.setdefault(key, len(attribute_sets))) for path, owner in zip(paths, owners, strict=True): if not path.steps: raise ValueError("PAT cannot encode an empty path") nodes, relations = node_rows[owner], relation_rows[owner] rows: list[int | tuple[str, int]] = [0] scales = [0.0] types = [int(PATTokenType.PATH_START)] hops = [0] constraint = constraints[owner] if constraint is not None: signal = float(constraint.max_hops - constraint.min_hops) if constraint.exact_hops is not None: signal += float(constraint.exact_hops) rows.append(0) scales.append(1.0 + signal / max_length) types.append(int(PATTokenType.CONSTRAINT)) hops.append(0) first = path.steps[0] rows += [node_offset + nodes[first.source_id], type_slot(first.source_type)] scales += [0.0, 0.0] types += [int(PATTokenType.ENTITY), int(PATTokenType.TYPE)] hops += [0, 0] for hop, step in enumerate(path.steps, start=1): relation = relations.get(step.relation_id) rows += [ 0 if relation is None else relation_offset + relation, attribute_slot(step.edge_attributes), node_offset + nodes[step.target_id], type_slot(step.target_type), 0, ] scales += [0.0, 0.0, 0.0, 0.0, float(hop) / max_length] types += [ int(PATTokenType.RELATION), int(PATTokenType.EDGE), int(PATTokenType.ENTITY), int(PATTokenType.TYPE), int(PATTokenType.HOP), ] hops += [hop] * 5 rows.append(0) scales.append(0.0) types.append(int(PATTokenType.PATH_END)) hops.append(len(path.steps)) token_rows.append(rows) token_scales.append(scales) token_types.append(types) token_hops.append(hops) device = query_states.device type_offset = relation_offset + relation_table.shape[0] attribute_offset = type_offset + len(type_ids) parts = [ query_states.new_zeros((1, query_states.shape[-1])), node_table.to(query_states.dtype), relation_table.to(query_states.dtype), ] if type_ids: parts.append( self.type_embedding( torch.tensor(list(type_ids), dtype=torch.long, device=device) ).to(query_states.dtype) ) if attribute_sets: flat = [index for key in attribute_sets for index in key] offsets = torch.tensor( [0, *itertools.accumulate(len(key) for key in attribute_sets)][:-1], dtype=torch.long, device=device, ) parts.append( functional.embedding_bag( torch.tensor(flat, dtype=torch.long, device=device), self.attribute_embedding.weight, offsets, mode="mean", ).to(query_states.dtype) ) table = torch.cat(parts) length = max(len(rows) for rows in token_rows) count = len(paths) index = torch.zeros((count, length), dtype=torch.long) scale = torch.zeros((count, length), dtype=query_states.dtype) padded_types = torch.full( (count, length), int(PATTokenType.PATH_END), dtype=torch.long ) padded_hops = torch.zeros((count, length), dtype=torch.long) padded_mask = torch.zeros((count, length), dtype=torch.long) for row, (rows, scales, types, hops) in enumerate( zip(token_rows, token_scales, token_types, token_hops, strict=True) ): size = len(rows) index[row, :size] = torch.tensor( [ slot if isinstance(slot, int) else (type_offset if slot[0] == "t" else attribute_offset) + slot[1] for slot in rows ], dtype=torch.long, ) scale[row, :size] = torch.tensor(scales, dtype=query_states.dtype) padded_types[row, :size] = torch.tensor(types, dtype=torch.long) padded_hops[row, :size] = torch.tensor(hops, dtype=torch.long) padded_mask[row, :size] = 1 index, scale = index.to(device), scale.to(device) owner_index = torch.tensor(list(owners), dtype=torch.long, device=device) features = table[index] + scale.unsqueeze(-1) * query_states[ owner_index ].unsqueeze(1) return PATInput( features, padded_types.to(device), padded_hops.to(device), padded_mask.to(device), ) def _type_state(self, type_name: str | None, query_state: Tensor) -> Tensor: index = stable_hash(type_name or "unknown", self.config.type_vocab_size) return self.type_embedding(torch.tensor(index, device=query_state.device)) def _attribute_state( self, attributes: Mapping[str, str], query_state: Tensor ) -> Tensor: if not attributes: return torch.zeros_like(query_state) indices = [ stable_hash(f"{key}={value}", self.config.attribute_vocab_size) for key, value in sorted(attributes.items()) ] ids = torch.tensor(indices, dtype=torch.long, device=query_state.device) return self.attribute_embedding(ids).mean(0) def encode_paths( self, paths: Sequence[GraphPath], node_states: Mapping[str, Tensor], relation_states: Mapping[str, Tensor], query_state: Tensor, constraint_state: ConstraintState | None = None, ) -> PATOutput: return self( self.build_input( paths, node_states, relation_states, query_state, constraint_state ), query_state if query_state.ndim == 2 else query_state.unsqueeze(0), )