File size: 2,442 Bytes
69928d6 | 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 | """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))
|