File size: 2,817 Bytes
e69b72a | 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 71 72 | """Causal segment boundaries and visibility masks."""
from __future__ import annotations
from dataclasses import dataclass
import torch
@dataclass(frozen=True, slots=True)
class SegmentLayout:
token_segment_ids: torch.Tensor
segment_token_mask: torch.Tensor
segment_valid: torch.Tensor
starts: torch.Tensor
ends: torch.Tensor
@property
def segment_count(self) -> int:
return int(self.segment_token_mask.shape[1])
def historical_memory_mask(self, memory_segment_ids: torch.Tensor) -> torch.Tensor:
"""Return [batch, token, memory] visibility for completed history only."""
if memory_segment_ids.ndim != 2:
raise ValueError("memory_segment_ids must have shape [batch, memory]")
token_segments = self.token_segment_ids.unsqueeze(-1)
return memory_segment_ids.unsqueeze(1) < token_segments
def fixed_segment_layout(
attention_mask: torch.Tensor,
*,
segment_size: int,
) -> SegmentLayout:
"""Create bounded contiguous segments without using future content.
Sentence and paragraph boundary commits can be supplied later as explicit
boundary metadata. Fixed maximum-length commits are always available and
preserve the same completed-history causality contract.
"""
if attention_mask.ndim != 2:
raise ValueError("attention_mask must have shape [batch, sequence]")
batch, seq_len = attention_mask.shape
device = attention_mask.device
segment_count = (seq_len + segment_size - 1) // segment_size
positions = torch.arange(seq_len, device=device)
token_segment_ids = torch.div(positions, segment_size, rounding_mode="floor")
token_segment_ids = token_segment_ids.unsqueeze(0).expand(batch, -1)
segment_ids = torch.arange(segment_count, device=device)
segment_token_mask = token_segment_ids.unsqueeze(1) == segment_ids.view(1, -1, 1)
segment_token_mask &= attention_mask.to(torch.bool).unsqueeze(1)
segment_valid = segment_token_mask.any(dim=-1)
starts = segment_ids * segment_size
ends = torch.minimum(starts + segment_size, torch.tensor(seq_len, device=device))
return SegmentLayout(
token_segment_ids=token_segment_ids,
segment_token_mask=segment_token_mask,
segment_valid=segment_valid,
starts=starts,
ends=ends,
)
def assert_no_open_segment_visibility(layout: SegmentLayout, memory_segment_ids: torch.Tensor) -> None:
visible = layout.historical_memory_mask(memory_segment_ids)
same_segment = memory_segment_ids.unsqueeze(1) == layout.token_segment_ids.unsqueeze(-1)
if bool((visible & same_segment).any()):
raise AssertionError("open segment can read its own compiled state")
__all__ = ["SegmentLayout", "assert_no_open_segment_visibility", "fixed_segment_layout"]
|