1import json
 2import numpy as np
 3import torch
 4from laplacian_attention import FirstSpikeLaplacianAttention
 5
 6
 7def main():
 8    torch.manual_seed(7)
 9    dtype = torch.float64
10    b, h, nq, nk, c, d = 2, 3, 5, 7, 4, 6
11    q = torch.rand(b, h, nq, c, dtype=dtype, requires_grad=True)
12    k = torch.rand(b, h, nk, c, dtype=dtype, requires_grad=True)
13    v = torch.rand(b, h, nk, d, dtype=dtype, requires_grad=True)
14    mod = FirstSpikeLaplacianAttention(h, init_sigma=.7).to(dtype)
15    out, a, _ = mod(q, k, v, return_attention=True)
16    out.square().mean().backward()
17
18    # 13 independent valid cases: near distance 0, far distance delta >= 0.
19    # The expected probability is the exact two-key normalization after L1.
20    delta = torch.linspace(0, 3, 13, dtype=dtype)
21    q2 = torch.zeros(13, 1, 1, 1, dtype=dtype)
22    k2 = torch.stack([torch.zeros_like(delta), delta], dim=-1).reshape(13, 1, 2, 1)
23    v2 = torch.eye(2, dtype=dtype).reshape(1, 1, 2, 2).expand(13, -1, -1, -1)
24    sigma = .7
25    mod2 = FirstSpikeLaplacianAttention(1, init_sigma=sigma).to(dtype)
26    _, a2, _ = mod2(q2, k2, v2, return_attention=True)
27    observed_prob = a2[:, 0, 0, 0].detach().numpy()
28    expected_prob = (1 / (1 + torch.exp(-delta / sigma))).numpy()
29
30    q3 = torch.zeros(1, 1, 1, 1, dtype=dtype)
31    k3 = torch.tensor([[[[.8], [2.4]]]], dtype=dtype)
32    _, a3, _ = mod2(q3, k3, torch.ones(1,1,2,1,dtype=dtype), return_attention=True)
33    observed_log_ratio = float(torch.log(a3[0,0,0,0] / a3[0,0,0,1]))
34    predicted_log_ratio = (2.4-.8)/sigma
35    result = {
36      'shape': list(out.shape),
37      'row_sum_error': float((a.sum(-1)-1).abs().max()),
38      'nonnegative': bool((a >= 0).all()),
39      'positive_bandwidths': mod.bandwidth().detach().tolist(),
40      'gradient_finite': bool(all(x.grad is not None and torch.isfinite(x.grad).all() for x in [q,k,v,mod.theta])),
41      'two_key_max_error': float(np.max(np.abs(observed_prob-expected_prob))),
42      'two_key_crossing_delta': float(delta[np.argmin(np.abs(observed_prob-.5))]),
43      'log_ratio_observed': observed_log_ratio,
44      'log_ratio_predicted': predicted_log_ratio,
45      'log_ratio_error': abs(observed_log_ratio-predicted_log_ratio),
46    }
47    print(json.dumps(result, indent=2))
48
49if __name__ == '__main__': main()