"""TensorViz teaching example: one untrained, unmasked attention head."""
import torch
import torch.nn as nn


class SelfAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.query = nn.Linear(32, 32)
        self.key = nn.Linear(32, 32)
        self.value = nn.Linear(32, 32)
        self.softmax = nn.Softmax(dim=-1)
        self.project = nn.Linear(32, 32)

    def forward(self, x):
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        scores = torch.matmul(q, k.transpose(-2, -1))
        scores = scores / (32 ** 0.5)
        weights = self.softmax(scores)
        mixed = torch.matmul(weights, v)
        return self.project(mixed) + x


# tensorviz: input=1,4,32
class TokenMixer(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = SelfAttention()
        self.norm = nn.LayerNorm(32)

    def forward(self, x):
        x = self.attention(x)
        return self.norm(x)


if __name__ == "__main__":
    torch.manual_seed(0)
    with torch.no_grad():
        output = TokenMixer().eval()(torch.ones(1, 4, 32))
    assert tuple(output.shape) == (1, 4, 32)
    print("Input [1, 4, 32] -> output", list(output.shape))