| |
| |
| """Shared evidence prefix and isolated candidate branch layout.""" |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
|
|
|
|
| @dataclass(frozen=True) |
| class CandidateTree: |
| prefix_length: int |
| branch_lengths: tuple[int, ...] |
|
|
| def __post_init__(self): |
| if self.prefix_length <= 0 or not self.branch_lengths or min(self.branch_lengths) <= 0: |
| raise ValueError("A nonempty prefix and nonempty branches are required") |
|
|
| @property |
| def length(self): |
| return self.prefix_length + sum(self.branch_lengths) |
|
|
| def allowed(self, *, device=None): |
| groups = [-1] * self.prefix_length |
| for index, length in enumerate(self.branch_lengths): |
| groups.extend([index] * length) |
| groups = torch.tensor(groups, device=device) |
| physical = torch.arange(self.length, device=device) |
| causal = physical[None, :] <= physical[:, None] |
| return causal & ((groups[None, :] == -1) | (groups[:, None] == groups[None, :])) |
|
|
| def additive_mask(self, *, dtype=torch.float32, device=None): |
| allowed = self.allowed(device=device) |
| return torch.zeros_like(allowed, dtype=dtype).masked_fill(~allowed, float("-inf"))[None, None] |
|
|
| def position_ids(self, prefix_positions): |
| if prefix_positions.ndim != 3 or prefix_positions.shape[-1] != self.prefix_length: |
| raise ValueError("Expected native prefix positions with shape (axes, batch, prefix)") |
| if prefix_positions.shape[1] != 1: |
| raise ValueError("This reference packs one request") |
| origin = int(prefix_positions.max()) + 1 |
| parts = [prefix_positions] |
| for length in self.branch_lengths: |
| branch = torch.arange(origin, origin + length, device=prefix_positions.device) |
| parts.append(branch.view(1, 1, -1).expand(prefix_positions.shape[0], 1, -1)) |
| return torch.cat(parts, dim=-1) |
|
|
| def attention_pairs(self): |
| prefix = self.prefix_length |
| return prefix * (prefix + 1) // 2 + sum( |
| prefix * length + length * (length + 1) // 2 for length in self.branch_lengths |
| ) |
|
|
