1"""First-Spike Laplacian attention for latency-coded tokens."""
 2import torch
 3from torch import nn
 4
 5
 6class FirstSpikeLaplacianAttention(nn.Module):
 7    """Attention using exp(-||q-k||_1 / sigma), with one bandwidth per head.
 8
 9    q, k: [batch, heads, query_tokens/key_tokens, channels]
10    v:    [batch, heads, key_tokens, value_channels]
11    """
12    def __init__(self, heads: int, init_sigma: float = 1.0, eps: float = 1e-6):
13        super().__init__()
14        if init_sigma <= eps:
15            raise ValueError("init_sigma must be positive")
16        self.eps = eps
17        # inverse softplus gives the requested initial positive bandwidth
18        theta = torch.log(torch.expm1(torch.tensor(float(init_sigma - eps))))
19        self.theta = nn.Parameter(theta.repeat(heads))
20
21    def bandwidth(self):
22        return torch.nn.functional.softplus(self.theta) + self.eps
23
24    def forward(self, q, k, v, return_attention=False):
25        if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
26            raise ValueError("q, k, v must have shape [B,H,T,C]")
27        if q.shape[0] != k.shape[0] or q.shape[1] != k.shape[1]:
28            raise ValueError("q and k batch/head dimensions must match")
29        if k.shape[:2] != v.shape[:2] or k.shape[2] != v.shape[2]:
30            raise ValueError("k and v batch/head/token dimensions must match")
31        distances = (q[:, :, :, None, :] - k[:, :, None, :, :]).abs().sum(dim=-1)
32        logits = -distances / self.bandwidth()[None, :, None, None]
33        attention = torch.softmax(logits, dim=-1)
34        output = attention @ v
35        return (output, attention, distances) if return_attention else output