1"""Gated local-global graph attention, single-head MVP."""
 2import torch
 3from torch import nn
 4
 5
 6def phi(x):
 7    return torch.nn.functional.elu(x) + 1.0
 8
 9
10def local_attention(x, edge_src, edge_dst, wq, wk, wv):
11    q, k, v = x @ wq, x @ wk, x @ wv
12    r = q.shape[-1]
13    logits = (q[edge_src] * k[edge_dst]).sum(-1) / (r ** 0.5)
14    n = x.shape[0]
15    mx = torch.full((n,), -torch.inf, device=x.device, dtype=x.dtype)
16    mx.scatter_reduce_(0, edge_src, logits, reduce="amax", include_self=True)
17    weights_un = torch.exp(logits - mx[edge_src])
18    den = torch.zeros(n, device=x.device, dtype=x.dtype)
19    den.scatter_add_(0, edge_src, weights_un)
20    weights = weights_un / den[edge_src].clamp_min(1e-12)
21    out = torch.zeros((n, v.shape[-1]), device=x.device, dtype=x.dtype)
22    out.index_add_(0, edge_src, weights[:, None] * v[edge_dst])
23    return out
24
25
26def global_linear_attention(x, wq, wk, wv, eps=1e-6):
27    q, k, v = x @ wq, x @ wk, x @ wv
28    pq, pk = phi(q), phi(k)
29    # S = sum_j outer(phi(k_j), v_j), z = sum_j phi(k_j)
30    s = pk.transpose(0, 1) @ v
31    z = pk.sum(0)
32    return (pq @ s) / ((pq @ z)[:, None] + eps)
33
34
35class GatedLocalGlobalAttention(nn.Module):
36    def __init__(self, d, r=None):
37        super().__init__()
38        r = r or d
39        self.wq, self.wk, self.wv = nn.Linear(d, r, bias=False), nn.Linear(d, r, bias=False), nn.Linear(d, r, bias=False)
40        self.out = nn.Linear(r, d)
41        self.gate = nn.Linear(d, 1)
42
43    def forward(self, x, edge_src, edge_dst, return_gate=False):
44        l = local_attention(x, edge_src, edge_dst, self.wq.weight.T, self.wk.weight.T, self.wv.weight.T)
45        g = global_linear_attention(x, self.wq.weight.T, self.wk.weight.T, self.wv.weight.T)
46        gate = torch.sigmoid(self.gate(x))
47        y = self.out(gate * l + (1.0 - gate) * g)
48        return (y, gate.squeeze(-1)) if return_gate else y
49
50
51def dense_attention(x, wq, wk, wv):
52    q, k, v = x @ wq, x @ wk, x @ wv
53    a = torch.softmax(q @ k.T / (q.shape[-1] ** 0.5), dim=-1)
54    return a @ v