source/model.py

from __future__ import annotations

import torch
from torch import nn
from torch.nn import functional as F


class FastWeightProgrammer(nn.Module):
    def __init__(
        self,
        key_vocab: int = 16,
        value_vocab: int = 11,
        values: int = 10,
        dimension: int = 64,
    ) -> None:
        super().__init__()
        self.key_embedding = nn.Embedding(key_vocab, dimension)
        self.value_embedding = nn.Embedding(value_vocab, dimension)
        self.write_controller = nn.Linear(2 * dimension + 1, 1)
        self.readout = nn.Linear(dimension, values)

    def forward(
        self,
        keys: torch.Tensor,
        values: torch.Tensor,
        writes: torch.Tensor,
        *,
        return_trace: bool = False,
    ):
        key_vectors = F.normalize(self.key_embedding(keys), dim=-1)
        value_vectors = self.value_embedding(values)
        controller_input = torch.cat(
            [key_vectors, value_vectors, writes.unsqueeze(-1)], dim=-1
        )
        strengths = torch.sigmoid(self.write_controller(controller_input))
        strengths = strengths.squeeze(-1) * writes
        memory = torch.einsum(
            "bt,btd,bte->bde",
            strengths,
            key_vectors,
            value_vectors,
        )
        query = key_vectors[:, -1]
        retrieved = torch.einsum("bd,bde->be", query, memory)
        logits = self.readout(retrieved)
        if return_trace:
            contributions = (
                torch.einsum("btd,bd->bt", key_vectors, query) * strengths
            )
            return logits, strengths, contributions
        return logits


class GRUControl(nn.Module):
    def __init__(
        self,
        key_vocab: int = 16,
        value_vocab: int = 11,
        values: int = 10,
        dimension: int = 16,
    ) -> None:
        super().__init__()
        self.key_embedding = nn.Embedding(key_vocab, dimension)
        self.value_embedding = nn.Embedding(value_vocab, dimension)
        self.recurrent = nn.GRU(2 * dimension + 1, dimension, batch_first=True)
        self.readout = nn.Linear(dimension, values)

    def forward(
        self, keys: torch.Tensor, values: torch.Tensor, writes: torch.Tensor
    ) -> torch.Tensor:
        inputs = torch.cat(
            [
                self.key_embedding(keys),
                self.value_embedding(values),
                writes.unsqueeze(-1),
            ],
            dim=-1,
        )
        hidden, _ = self.recurrent(inputs)
        return self.readout(hidden[:, -1])


def parameter_count(module: nn.Module) -> int:
    return sum(parameter.numel() for parameter in module.parameters())