File size: 1,871 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 | """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
|