# "Attention Is a Hopfield Update": the Hopfield layer of the hopfield-layers package (Ramsauer et al., 2021) with one
# update step and its input normalisation switched off computes exactly PyTorch's multi-head attention.
#   pip install git+https://github.com/ml-jku/hopfield-layers      (or pass the path of a local clone as argument)
import sys, torch, torch.nn as nn
if len(sys.argv) > 1: sys.path.insert(0, sys.argv[1])
from hflayers import Hopfield
torch.manual_seed(0); torch.set_grad_enabled(False)
d, heads, n, b = 32, 4, 10, 3
off = dict(normalize_stored_pattern=False, normalize_stored_pattern_affine=False, normalize_state_pattern=False,
           normalize_state_pattern_affine=False, normalize_pattern_projection=False, normalize_pattern_projection_affine=False)
hop = Hopfield(input_size=d, num_heads=heads, **off)              # update_steps_max = 0: exactly one update
mha = nn.MultiheadAttention(d, heads, batch_first=True)
core = hop.association_core
mha.in_proj_weight.copy_(core.in_proj_weight); mha.in_proj_bias.copy_(core.in_proj_bias)
mha.out_proj.weight.copy_(core.out_proj.weight); mha.out_proj.bias.copy_(core.out_proj.bias)
x = torch.randn(b, n, d)
print("self-attention, max |difference|:", (hop(x) - mha(x, x, x, need_weights=False)[0]).abs().max().item())
q, kv = torch.randn(b, 4, d), torch.randn(b, n, d)                # cross-attention: queries from one set, memories from another
print("cross-attention, max |difference|:", (hop((kv, q, kv)) - mha(q, kv, kv, need_weights=False)[0]).abs().max().item())
# with the default input LayerNorms switched on, the two are no longer identical
hop2 = Hopfield(input_size=d, num_heads=heads); hop2.association_core.load_state_dict(core.state_dict())
print("with the default LayerNorms on, max |difference|:", (hop2(x) - mha(x, x, x, need_weights=False)[0]).abs().max().item())