| """Constraint-membership embeddings with no board geometry assumptions.""" |
|
|
| from collections.abc import Sequence |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| def membership_embeddings( |
| tables: Sequence[nn.Module], memberships: Sequence[torch.Tensor] |
| ) -> torch.Tensor: |
| """Sum family embeddings in their declared order (also preserves legacy order).""" |
| if not tables or len(tables) != len(memberships): |
| raise ValueError("Provide one membership index tensor per embedding table") |
| result = tables[0](memberships[0]) |
| for table, indices in zip(tables[1:], memberships[1:], strict=True): |
| result = result + table(indices) |
| return result |
|
|
|
|
| class GroupMembershipEmbedding(nn.Module): |
| """Embed static memberships; input shape is (families, positions).""" |
|
|
| memberships: torch.Tensor |
|
|
| def __init__( |
| self, group_counts: Sequence[int], memberships: torch.Tensor, d_model: int |
| ) -> None: |
| super().__init__() |
| if memberships.ndim != 2 or len(group_counts) != memberships.shape[0]: |
| raise ValueError("Memberships must have shape (families, positions)") |
| if not group_counts or memberships.dtype != torch.long: |
| raise ValueError("At least one family with integer memberships is required") |
| for count, indices in zip(group_counts, memberships, strict=True): |
| if count <= 0 or torch.any((indices < 0) | (indices >= count)): |
| raise ValueError("Membership index outside its family") |
| self.tables = nn.ModuleList([nn.Embedding(count, d_model) for count in group_counts]) |
| self.register_buffer("memberships", memberships.clone()) |
|
|
| def forward(self) -> torch.Tensor: |
| return membership_embeddings(list(self.tables), list(self.memberships.unbind())) |
|
|
