| """Sparse typed relation operators for predicate-hypergraph traversal.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Mapping, Sequence |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class SparseRelationOperators(nn.Module): |
| """A fixed set of sparse source-to-target adjacency operators. |
| |
| Every edge is ``(relation, source, target, weight)``. A traversal applies |
| ``A_relation.T @ state``. Relation names choose distinct operators rather |
| than modulating a shared learned message function. |
| """ |
|
|
| def __init__( |
| self, |
| node_count: int, |
| relation_names: Sequence[str], |
| edges: Mapping[str, Sequence[tuple[int, int] | tuple[int, int, float]]], |
| ) -> None: |
| super().__init__() |
| if node_count <= 0: |
| raise ValueError("node_count must be positive") |
| names = tuple(str(name) for name in relation_names) |
| if not names or len(set(names)) != len(names): |
| raise ValueError("relation_names must be non-empty and unique") |
| relation_ids: list[int] = [] |
| sources: list[int] = [] |
| targets: list[int] = [] |
| weights: list[float] = [] |
| for relation_id, name in enumerate(names): |
| for edge in edges.get(name, ()): |
| if len(edge) == 2: |
| source, target = edge |
| weight = 1.0 |
| elif len(edge) == 3: |
| source, target, weight = edge |
| else: |
| raise ValueError("edges must be (source, target) or (source, target, weight)") |
| source = int(source) |
| target = int(target) |
| if not 0 <= source < node_count or not 0 <= target < node_count: |
| raise IndexError(f"edge ({source}, {target}) is outside {node_count} nodes") |
| if float(weight) < 0: |
| raise ValueError("relation weights must be non-negative") |
| relation_ids.append(relation_id) |
| sources.append(source) |
| targets.append(target) |
| weights.append(float(weight)) |
| self.node_count = int(node_count) |
| self.relation_names = names |
| self.register_buffer("edge_relations", torch.tensor(relation_ids, dtype=torch.long)) |
| self.register_buffer("edge_sources", torch.tensor(sources, dtype=torch.long)) |
| self.register_buffer("edge_targets", torch.tensor(targets, dtype=torch.long)) |
| self.register_buffer("edge_weights", torch.tensor(weights, dtype=torch.float32)) |
|
|
| @property |
| def relation_count(self) -> int: |
| return len(self.relation_names) |
|
|
| def relation_id(self, relation: str | int) -> int: |
| if isinstance(relation, int): |
| if not 0 <= relation < self.relation_count: |
| raise IndexError("relation index out of range") |
| return relation |
| try: |
| return self.relation_names.index(relation) |
| except ValueError as exc: |
| raise KeyError(f"unknown relation {relation!r}") from exc |
|
|
| def sparse_operator(self, relation: str | int, *, transpose: bool = False) -> torch.Tensor: |
| relation_id = self.relation_id(relation) |
| selected = self.edge_relations == relation_id |
| source = self.edge_sources[selected] |
| target = self.edge_targets[selected] |
| if transpose: |
| source, target = target, source |
| indices = torch.stack([source, target]) if source.numel() else torch.empty( |
| 2, 0, dtype=torch.long, device=self.edge_sources.device |
| ) |
| values = self.edge_weights[selected] |
| return torch.sparse_coo_tensor( |
| indices, |
| values, |
| (self.node_count, self.node_count), |
| device=values.device, |
| ).coalesce() |
|
|
| def step( |
| self, |
| state: torch.Tensor, |
| relation: str | int, |
| *, |
| semiring: str = "sum_product", |
| ) -> torch.Tensor: |
| if state.shape[-1] != self.node_count: |
| raise ValueError(f"state must end with {self.node_count} nodes") |
| relation_id = self.relation_id(relation) |
| selected = self.edge_relations == relation_id |
| source = self.edge_sources[selected] |
| target = self.edge_targets[selected] |
| weight = self.edge_weights[selected].to(dtype=state.dtype) |
| flat = state.reshape(-1, self.node_count) |
| if semiring == "sum_product": |
| out = flat.new_zeros(flat.shape) |
| if source.numel(): |
| contribution = flat[:, source] * weight.unsqueeze(0) |
| out.scatter_add_(1, target.unsqueeze(0).expand(flat.shape[0], -1), contribution) |
| elif semiring == "max_product": |
| out = flat.new_zeros(flat.shape) |
| if source.numel(): |
| contribution = flat[:, source] * weight.unsqueeze(0) |
| out.scatter_reduce_( |
| 1, |
| target.unsqueeze(0).expand(flat.shape[0], -1), |
| contribution, |
| reduce="amax", |
| include_self=True, |
| ) |
| elif semiring == "logsumexp": |
| out = flat.new_full(flat.shape, -torch.inf) |
| if source.numel(): |
| log_weight = weight.clamp_min(torch.finfo(weight.dtype).tiny).log() |
| contribution = flat[:, source] + log_weight.unsqueeze(0) |
| index = target.unsqueeze(0).expand(flat.shape[0], -1) |
| maximum = flat.new_full(flat.shape, -torch.inf) |
| maximum.scatter_reduce_(1, index, contribution, reduce="amax", include_self=True) |
| gathered_max = maximum.gather(1, index) |
| shifted = torch.where( |
| torch.isfinite(gathered_max), |
| (contribution - gathered_max).exp(), |
| torch.zeros_like(contribution), |
| ) |
| total = flat.new_zeros(flat.shape) |
| total.scatter_add_(1, index, shifted) |
| finite = torch.isfinite(maximum) & (total > 0) |
| out = torch.where(finite, maximum + total.clamp_min(torch.finfo(total.dtype).tiny).log(), out) |
| else: |
| raise ValueError("semiring must be sum_product, max_product, or logsumexp") |
| return out.reshape(state.shape) |
|
|
|
|
| __all__ = ["SparseRelationOperators"] |
|
|
