model.py

from __future__ import annotations

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


class DenseGraphAttentionLayer(nn.Module):
    def __init__(self, input_dimensions: int, output_dimensions: int, heads: int) -> None:
        super().__init__()
        self.output_dimensions = output_dimensions
        self.heads = heads
        self.weight = nn.Parameter(
            torch.randn(heads, input_dimensions, output_dimensions) * 0.1
        )
        self.source_attention = nn.Parameter(torch.randn(heads, output_dimensions) * 0.1)
        self.target_attention = nn.Parameter(torch.randn(heads, output_dimensions) * 0.1)

    def forward(
        self,
        features: torch.Tensor,
        adjacency: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        projected = torch.einsum("ni,hio->hno", features, self.weight)
        source = torch.einsum("hno,ho->hn", projected, self.source_attention)
        target = torch.einsum("hno,ho->hn", projected, self.target_attention)
        scores = F.leaky_relu(source[:, :, None] + target[:, None, :], 0.2)
        mask = adjacency.bool() | torch.eye(len(adjacency), device=adjacency.device).bool()
        scores = scores.masked_fill(~mask[None], -torch.inf)
        attention = torch.softmax(scores, dim=2)
        output = torch.einsum("hij,hjo->hio", attention, projected)
        return output.permute(1, 0, 2).reshape(len(features), -1), attention


class MeshGraphGAT(nn.Module):
    def __init__(self, input_dimensions: int = 8) -> None:
        super().__init__()
        self.first = DenseGraphAttentionLayer(input_dimensions, 16, heads=4)
        self.second = DenseGraphAttentionLayer(64, 2, heads=1)

    def forward(
        self,
        features: torch.Tensor,
        adjacency: torch.Tensor,
        *,
        return_attention: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        hidden, first_attention = self.first(features, adjacency)
        hidden = F.elu(hidden)
        hidden = F.dropout(hidden, p=0.15, training=self.training)
        logits, _ = self.second(hidden, adjacency)
        if return_attention:
            return logits, first_attention
        return logits


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