| """Explicit evidence policies for soft symbolic distributions.""" |
|
|
| from typing import Literal |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| EvidenceMode = Literal["observed", "predicted", "none"] |
|
|
|
|
| def clamp_evidence( |
| logits: torch.Tensor, |
| mask: torch.Tensor | None, |
| *, |
| mode: EvidenceMode, |
| observed_symbols: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Produce soft symbols, optionally anchoring observed positions. |
| |
| Observed symbols are an explicit training input, never inferred from labels |
| hidden in a batch. The caller selects the policy for its instantiation. |
| """ |
| probabilities = F.softmax(logits, dim=-1) |
| if mode == "none": |
| return probabilities |
| if mode not in ("observed", "predicted"): |
| raise ValueError(f"Unknown evidence mode: {mode}") |
| if mask is None or mask.shape != logits.shape[:-1]: |
| raise ValueError("Evidence mask must match the symbolic positions") |
| if mode == "observed": |
| if observed_symbols is None or observed_symbols.shape != mask.shape: |
| raise ValueError("Observed evidence requires symbols matching the mask") |
| symbols = observed_symbols |
| else: |
| symbols = logits.argmax(dim=-1) |
| weights = mask.unsqueeze(-1).float() |
| anchors = F.one_hot(symbols, num_classes=logits.shape[-1]).float() |
| return weights * anchors + (1.0 - weights) * probabilities |
|
|
|
|
| def restore_evidence( |
| post_logits: torch.Tensor, pre_logits: torch.Tensor, mask: torch.Tensor |
| ) -> torch.Tensor: |
| """Optional token-head policy; not appropriate for a pooled task head.""" |
| if post_logits.shape != pre_logits.shape or mask.shape != pre_logits.shape[:-1]: |
| raise ValueError("Evidence restoration requires matching token logits and mask") |
| weights = mask.unsqueeze(-1).float() |
| return weights * pre_logits + (1.0 - weights) * post_logits |
|
|
