torch.nn.AssocEmb

assoc_emb.py

import torch


class AssocEmb(torch.nn.Module):
    def __init__(self, dim: list[int]):
        super(AssocEmb, self).__init__()
        self.dim = dim
        self.tbl = torch.nn.ParameterDict()

    def forward(self, idx: str) -> torch.Tensor:
        if idx not in self.tbl:
            p = torch.rand(self.dim, requires_grad=True)
            self.tbl[idx] = torch.nn.Parameter(p)
        return self.tbl[idx]


m = AssocEmb([2, 2])
x = m("apple")
y = m("orange")
z = x * y
c = z.sum()
c.backward()
for k, v in enumerate(m.tbl):
    print(v, m.tbl[v], m.tbl[v].grad)