| """The exclusivity operator over explicit groups, independent of board geometry.""" |
|
|
| from typing import NamedTuple |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class OperatorEvaluation(NamedTuple): |
| probabilities: torch.Tensor |
| consequences: torch.Tensor |
|
|
|
|
| def exclusivity_consequence( |
| probabilities: torch.Tensor, |
| log_complement: torch.Tensor, |
| member_dims: int | tuple[int, ...], |
| ) -> torch.Tensor: |
| """Exclusion product on group-member axes of a supplied tensor layout. |
| |
| Index-based groups and structured partition views share this kernel. Views |
| let compatibility adapters preserve reduction order and autograd accumulation. |
| """ |
| excluded = torch.exp(log_complement.sum(dim=member_dims, keepdim=True) - log_complement) |
| return probabilities * excluded |
|
|
|
|
| class GroupExclusivityOperator(nn.Module): |
| """One rectangular family of groups; different families may have different sizes. |
| |
| This is the exclusivity family of the probabilistic lift, not an interpreter |
| for arbitrary ASP programs. Indices are derived configuration, not weights. |
| """ |
|
|
| groups: torch.Tensor |
|
|
| def __init__(self, groups: torch.Tensor, num_positions: int, eps: float = 1e-7) -> None: |
| super().__init__() |
| if groups.ndim != 2 or groups.numel() == 0 or groups.dtype != torch.long: |
| raise ValueError("Groups must be a nonempty integer (groups, members) tensor") |
| if num_positions < 1 or torch.any((groups < 0) | (groups >= num_positions)): |
| raise ValueError("Group member outside the declared position range") |
| if any(len(set(row)) != len(row) for row in groups.tolist()): |
| raise ValueError("A group cannot contain the same position twice") |
| if not 0 < eps < 1: |
| raise ValueError("eps must lie strictly between zero and one") |
| self.num_positions = num_positions |
| self.eps = eps |
| self.register_buffer("groups", groups.clone(), persistent=False) |
|
|
| def forward(self, probabilities: torch.Tensor) -> OperatorEvaluation: |
| if probabilities.ndim != 3 or probabilities.shape[1] != self.num_positions: |
| raise ValueError("Expected probabilities shaped (batch, declared positions, symbols)") |
| grouped = probabilities[:, self.groups, :] |
| log_complement = torch.log1p(-grouped.clamp(max=1.0 - self.eps)) |
| return OperatorEvaluation(grouped, exclusivity_consequence(grouped, log_complement, -2)) |
|
|
