| """Exact tensor-product role binding with a fixed orthonormal role basis.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| DEFAULT_ROLE_NAMES = ( |
| "ARG0", |
| "ARG1", |
| "ARG2", |
| "TIME", |
| "LOCATION", |
| "OTHER", |
| "NULL", |
| ) |
|
|
|
|
| class OrthonormalRoleBinder(nn.Module): |
| """Bind filler vectors to canonical role vectors and unbind them exactly. |
| |
| The memory matrix is ``sum_r filler[r] outer role[r]``. Canonical one-hot |
| role vectors avoid numerical drift and make the Phase 0 contract exact. |
| A supplied orthonormal basis can be used for controlled permutation tests. |
| """ |
|
|
| def __init__( |
| self, |
| filler_dim: int, |
| role_names: Sequence[str] = DEFAULT_ROLE_NAMES, |
| *, |
| role_dim: int | None = None, |
| basis: torch.Tensor | None = None, |
| ) -> None: |
| super().__init__() |
| if filler_dim <= 0: |
| raise ValueError("filler_dim must be positive") |
| names = tuple(str(name) for name in role_names) |
| if not names or len(set(names)) != len(names): |
| raise ValueError("role_names must be non-empty and unique") |
| role_dim = len(names) if role_dim is None else int(role_dim) |
| if role_dim < len(names): |
| raise ValueError("role_dim must be at least the number of roles") |
| if basis is None: |
| basis = torch.eye(role_dim, dtype=torch.float32)[: len(names)] |
| else: |
| basis = basis.detach().to(dtype=torch.float32) |
| if basis.shape != (len(names), role_dim): |
| raise ValueError( |
| f"basis must have shape {(len(names), role_dim)}, got {tuple(basis.shape)}" |
| ) |
| gram = basis @ basis.transpose(0, 1) |
| identity = torch.eye(len(names), dtype=gram.dtype, device=gram.device) |
| if not torch.allclose(gram, identity, atol=1e-6, rtol=1e-6): |
| raise ValueError("role basis must be orthonormal") |
| self.filler_dim = int(filler_dim) |
| self.role_names = names |
| self.role_dim = role_dim |
| self.register_buffer("basis", basis.contiguous(), persistent=True) |
|
|
| @property |
| def role_count(self) -> int: |
| return len(self.role_names) |
|
|
| def role_id(self, name: str) -> int: |
| try: |
| return self.role_names.index(name) |
| except ValueError as exc: |
| raise KeyError(f"unknown role {name!r}") from exc |
|
|
| def role_vector(self, roles: int | str | torch.Tensor) -> torch.Tensor: |
| if isinstance(roles, str): |
| roles = self.role_id(roles) |
| if isinstance(roles, int): |
| if not 0 <= roles < self.role_count: |
| raise IndexError("role index out of range") |
| return self.basis[roles] |
| role_ids = roles.to(device=self.basis.device, dtype=torch.long) |
| if bool(((role_ids < 0) | (role_ids >= self.role_count)).any()): |
| raise IndexError("role index out of range") |
| return self.basis[role_ids] |
|
|
| def bind(self, fillers: torch.Tensor, *, role_mask: torch.Tensor | None = None) -> torch.Tensor: |
| """Return a TPR matrix from ``[..., roles, filler_dim]`` fillers.""" |
| if fillers.shape[-2:] != (self.role_count, self.filler_dim): |
| raise ValueError( |
| "fillers must end with " |
| f"[{self.role_count}, {self.filler_dim}], got {tuple(fillers.shape)}" |
| ) |
| values = fillers |
| if role_mask is not None: |
| expected = fillers.shape[:-1] |
| if role_mask.shape != expected: |
| raise ValueError(f"role_mask must have shape {expected}, got {tuple(role_mask.shape)}") |
| values = values * role_mask.to(device=values.device, dtype=values.dtype).unsqueeze(-1) |
| basis = self.basis.to(device=values.device, dtype=values.dtype) |
| return torch.einsum("...rf,rd->...fd", values, basis) |
|
|
| def unbind(self, memory: torch.Tensor, roles: int | str | torch.Tensor) -> torch.Tensor: |
| """Apply the requested role as an algebraic unbinding operator.""" |
| if memory.shape[-2:] != (self.filler_dim, self.role_dim): |
| raise ValueError( |
| "memory must end with " |
| f"[{self.filler_dim}, {self.role_dim}], got {tuple(memory.shape)}" |
| ) |
| role_vector = self.role_vector(roles).to(device=memory.device, dtype=memory.dtype) |
| if role_vector.ndim == 1: |
| return torch.einsum("...fd,d->...f", memory, role_vector) |
| if role_vector.shape[:-1] != memory.shape[:-2]: |
| raise ValueError( |
| "tensor role IDs must match memory batch dimensions: " |
| f"{tuple(role_vector.shape[:-1])} != {tuple(memory.shape[:-2])}" |
| ) |
| return torch.einsum("...fd,...d->...f", memory, role_vector) |
|
|
| def unbind_all(self, memory: torch.Tensor) -> torch.Tensor: |
| if memory.shape[-2:] != (self.filler_dim, self.role_dim): |
| raise ValueError("memory has incompatible trailing dimensions") |
| basis = self.basis.to(device=memory.device, dtype=memory.dtype) |
| return torch.einsum("...fd,rd->...rf", memory, basis) |
|
|
| def permuted(self, permutation: torch.Tensor) -> "OrthonormalRoleBinder": |
| """Return a binder with globally relabeled role basis vectors.""" |
| permutation = permutation.to(dtype=torch.long, device=self.basis.device) |
| if permutation.shape != (self.role_count,): |
| raise ValueError(f"permutation must have shape ({self.role_count},)") |
| if sorted(permutation.tolist()) != list(range(self.role_count)): |
| raise ValueError("permutation must contain each role exactly once") |
| return OrthonormalRoleBinder( |
| self.filler_dim, |
| self.role_names, |
| role_dim=self.role_dim, |
| basis=self.basis[permutation].cpu(), |
| ) |
|
|
|
|
| __all__ = ["DEFAULT_ROLE_NAMES", "OrthonormalRoleBinder"] |
|
|
