| """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"] |
|
|
