File size: 1,780 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 | """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()))
|