File size: 1,677 Bytes
ce680a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | # Auto-extracted class source (static)
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) |