| """Compile the predicted three-stream state consumed by the LM path.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
|
|
| from strata.modeling.ph_pat.config import PHPATArm, PHPATConfig |
| from strata.modeling.ph_pat.dependency_chart import DependencyChart, DependencyChartOutput |
| from strata.modeling.ph_pat.predicate_hypergraph import PredicateHypergraphCompiler, PredicateHypergraphState |
| from strata.modeling.ph_pat.primitive_registers import PrimitiveRegisterState |
| from strata.modeling.ph_pat.role_primitive_bridge import BridgeOutput, RolePrimitiveBridge |
| from strata.modeling.ph_pat.segment_commit import SegmentLayout |
| from strata.modeling.ph_pat.span_compiler import SpanCompiler, SpanCompilerOutput |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class CompiledPHPATChart: |
| spans: SpanCompilerOutput |
| dependency: DependencyChartOutput |
| hypergraph: PredicateHypergraphState |
| primitives: PrimitiveRegisterState | None |
| bridge_permutation: torch.Tensor |
| arm: str |
|
|
|
|
| class ChartExecutor(nn.Module): |
| def __init__(self, config: PHPATConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.span_compiler = SpanCompiler(config) |
| self.dependency_chart = DependencyChart(config) |
| self.hypergraph = PredicateHypergraphCompiler(config) |
| self.bridge = RolePrimitiveBridge(config) |
|
|
| def forward( |
| self, |
| hidden: torch.Tensor, |
| layout: SegmentLayout, |
| *, |
| arm: PHPATArm | None = None, |
| corruption: str = "none", |
| generator: torch.Generator | None = None, |
| ) -> CompiledPHPATChart: |
| selected = arm or self.config.arm |
| spans = self.span_compiler(hidden, layout) |
| dependency = self.dependency_chart(spans) |
| typed_event = selected in {PHPATArm.TYPED_EVENT, PHPATArm.FULL, PHPATArm.SHUFFLED_BRIDGE, PHPATArm.RANDOM_CHART} |
| typed_primitive = selected in {PHPATArm.TYPED_PRIMITIVE, PHPATArm.FULL, PHPATArm.SHUFFLED_BRIDGE, PHPATArm.RANDOM_CHART} |
| hypergraph = self.hypergraph(hidden, layout, dependency, typed_roles=typed_event) |
|
|
| if selected == PHPATArm.GENERIC_MEMORY: |
| pooled = hypergraph.frames.mean(dim=2, keepdim=True) |
| hypergraph = hypergraph.with_frames(pooled.expand_as(hypergraph.frames)) |
| elif selected == PHPATArm.RANDOM_CHART or corruption == "random_matched": |
| frames = _online_random_matched(hypergraph.frames, generator=generator) |
| hypergraph = hypergraph.with_frames(frames) |
|
|
| use_primitives = selected in { |
| PHPATArm.TYPED_PRIMITIVE, |
| PHPATArm.FULL, |
| PHPATArm.RANDOM_CHART, |
| PHPATArm.SHUFFLED_BRIDGE, |
| } |
| bridge: BridgeOutput | None = None |
| if use_primitives: |
| bridge = self.bridge( |
| hypergraph, |
| typed_primitives=typed_primitive, |
| shuffle_bridge=selected == PHPATArm.SHUFFLED_BRIDGE or corruption == "bridge_shuffle", |
| generator=generator, |
| ) |
| permutation = ( |
| bridge.role_permutation |
| if bridge is not None |
| else torch.arange(self.config.role_count, device=hidden.device) |
| ) |
| return CompiledPHPATChart( |
| spans=spans, |
| dependency=dependency, |
| hypergraph=hypergraph, |
| primitives=bridge.primitives if bridge is not None else None, |
| bridge_permutation=permutation, |
| arm=selected.value, |
| ) |
|
|
|
|
| def _online_random_matched(tensor: torch.Tensor, *, generator: torch.Generator | None) -> torch.Tensor: |
| sample = torch.randn(tensor.shape, device=tensor.device, dtype=torch.float32, generator=generator) |
| source = tensor.float() |
| reduce_dims = tuple(range(1, tensor.ndim)) |
| mean = source.mean(dim=reduce_dims, keepdim=True) |
| std = source.std(dim=reduce_dims, keepdim=True).clamp_min(1e-6) |
| sample = (sample - sample.mean(dim=reduce_dims, keepdim=True)) / sample.std( |
| dim=reduce_dims, keepdim=True |
| ).clamp_min(1e-6) |
| return (sample * std + mean).to(tensor.dtype) |
|
|
|
|
| __all__ = ["ChartExecutor", "CompiledPHPATChart"] |
|
|
