| |
|
|
| class DiracGraphConv(nn.Module): |
| def __init__(self, in_dim: int, out_dim: int, alpha: float = 1.0, bias: bool = True): |
| super().__init__() |
| self.lin = nn.Linear(in_dim, out_dim, bias=bias) |
| self.alpha = nn.Parameter(torch.tensor(alpha, dtype=torch.float32)) |
| self.bias_edge = nn.Parameter(torch.tensor(0.0, dtype=torch.float32)) |
|
|
| @staticmethod |
| def cosine_corr(z_i: torch.Tensor, z_j: torch.Tensor, eps: float = 1e-9) -> torch.Tensor: |
| num = (z_i * z_j).sum(dim=-1) |
| den = torch.clamp(z_i.norm(dim=-1) * z_j.norm(dim=-1), min=eps) |
| return num / den |
|
|
| def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor: |
| N = x.size(0) |
| row, col = edge_index |
| corr = self.cosine_corr(z[row], z[col]) |
| logits = self.alpha * corr + self.bias_edge |
| device = x.device |
| E = row.size(0) |
| ones = torch.ones(E, device=device) |
| max_per_row = torch.full((N,), -1e9, device=device) |
| max_per_row = max_per_row.index_put((row,), logits, accumulate=False).scatter_reduce_(0, row, logits, reduce="amax") |
| logits_centered = logits - max_per_row[row] |
| exp_logits = torch.exp(logits_centered) |
| denom = torch.zeros(N, device=device).index_add_(0, row, exp_logits) |
| attn = exp_logits / (denom[row] + 1e-9) |
| deg = torch.zeros(N, device=device).index_add_(0, row, ones) |
| norm = 1.0 / torch.clamp(deg[row], min=1.0) |
| msgs = norm.unsqueeze(-1) * attn.unsqueeze(-1) * x[col] |
| out = torch.zeros_like(x).index_add_(0, row, msgs) |
| return self.lin(out) |
