import torchimport torch.nn as nnimport torch.nn.functional as F
class GroupNorm(nn.GroupNorm):    """Group Normalization with 1 group.    Input: tensor in shape [B, C, H, W]    """    def __init__(self, num_channels, **kwargs):        super().__init__(1, num_channels, **kwargs)
class Conv2d_BN(nn.Sequential):    def __init__(self, a, b, ks=1, stride=1, pad=0, dilation=1,                 groups=1, bn_weight_init=1):        super().__init__()        self.add_module('c', nn.Conv2d(            a, b, ks, stride, pad, dilation, groups, bias=False))        self.add_module('bn', nn.BatchNorm2d(b))        nn.init.constant_(self.bn.weight, bn_weight_init)        nn.init.constant_(self.bn.bias, 0)
    @torch.no_grad()    def fuse(self):        c, bn = self._modules.values()        w = bn.weight / (bn.running_var + bn.eps)**0.5        b = bn.bias - self.bn.running_mean * \            self.bn.weight / (bn.running_var + bn.eps)**0.5        m = nn.Conv2d(w.size(1) * self.c.groups, w.size(            0), w.shape[2:], stride=self.c.stride, padding=self.c.padding, dilation=self.c.dilation, groups=self.c.groups,            device=c.weight.device)        m.weight.data.copy_(w)        m.bias.data.copy_(b)        return m
class BN_Linear(nn.Sequential):    def __init__(self, a, b, bias=True, std=0.02):        super().__init__()        self.add_module('bn', nn.BatchNorm1d(a))        self.add_module('l', nn.Linear(a, b, bias=bias))        trunc_normal_(self.l.weight, std=std)        if bias:            nn.init.constant_(self.l.bias, 0)
    @torch.no_grad()    def fuse(self):        bn, l = self._modules.values()        w = bn.weight / (bn.running_var + bn.eps)**0.5        b = bn.bias - self.bn.running_mean * \            self.bn.weight / (bn.running_var + bn.eps)**0.5        b = (l.weight @ b[:, None]).view(-1) + self.l.bias        m = nn.Linear(w.size(1), w.size(0))        m.weight.data.copy_(w)        m.bias.data.copy_(b)        return m
class SHSA(nn.Module):    """Single-Head Self-Attention"""    def __init__(self, dim, qk_dim, pdim):        super().__init__()        self.scale = qk_dim ** -0.5        self.qk_dim = qk_dim        self.dim = dim        self.pdim = pdim        self.pre_norm = GroupNorm(pdim)        self.qkv = Conv2d_BN(pdim, qk_dim * 2 + pdim)        self.proj = nn.Sequential(nn.ReLU(), Conv2d_BN(            dim, dim, bn_weight_init=0))
    def forward(self, x):        B, C, H, W = x.shape        x1, x2 = torch.split(x, [self.pdim, self.dim - self.pdim], dim=1)        x1 = self.pre_norm(x1)        qkv = self.qkv(x1)        q, k, v = qkv.split([self.qk_dim, self.qk_dim, self.pdim], dim=1)        q, k, v = q.flatten(2), k.flatten(2), v.flatten(2)        attn = (q.transpose(-2, -1) @ k) * self.scale        attn = attn.softmax(dim=-1)        x1 = (v @ attn.transpose(-2, -1)).reshape(B, self.pdim, H, W)        x = self.proj(torch.cat([x1, x2], dim=1))        return x
if __name__ == '__main__':    x = torch.randn(1, 64, 32, 32)    shsa = SHSA(dim=64, qk_dim=64, pdim=64)    print(shsa)    output = shsa(x)    print(f"Input shape: {x.shape}")    print(f"Output shape: {output.shape}")