File size: 2,273 Bytes
7c5e40e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Short typed graph programs with exact role unbinding."""
from __future__ import annotations
from dataclasses import dataclass
import torch
from strata.modeling.algebra.relations import SparseRelationOperators
from strata.modeling.algebra.role_binding import OrthonormalRoleBinder
@dataclass(frozen=True, slots=True)
class GraphProgram:
relations: tuple[str, ...] = ()
read_role: str | None = None
def __post_init__(self) -> None:
if len(self.relations) > 3:
raise ValueError("qualification graph programs are bounded to three hops")
@dataclass(frozen=True, slots=True)
class GraphProgramResult:
node_state: torch.Tensor
filler: torch.Tensor | None
trace: tuple[torch.Tensor, ...]
class GraphProgramExecutor:
def __init__(
self,
relations: SparseRelationOperators,
role_binder: OrthonormalRoleBinder,
) -> None:
self.relations = relations
self.role_binder = role_binder
def execute(
self,
anchor_state: torch.Tensor,
program: GraphProgram,
*,
event_memories: torch.Tensor | None = None,
semiring: str = "max_product",
) -> GraphProgramResult:
state = anchor_state
trace = [state]
for relation in program.relations:
state = self.relations.step(state, relation, semiring=semiring)
trace.append(state)
filler = None
if program.read_role is not None:
if event_memories is None:
raise ValueError("event_memories are required for a role read")
expected = (
*state.shape[:-1],
self.relations.node_count,
self.role_binder.filler_dim,
self.role_binder.role_dim,
)
if event_memories.shape != expected:
raise ValueError(f"event_memories must have shape {expected}, got {tuple(event_memories.shape)}")
memory = torch.einsum("...n,...nfd->...fd", state, event_memories)
filler = self.role_binder.unbind(memory, program.read_role)
return GraphProgramResult(node_state=state, filler=filler, trace=tuple(trace))
__all__ = ["GraphProgram", "GraphProgramExecutor", "GraphProgramResult"]
|