import torch
import torch.nn as nn
import matterhorn_pytorch.snn as snn
import torch.nn.functional as F
from rich import print
from typing import Callable


def main():
    m = snn.TemporalWiseAttention(4, 0.5)
    i = torch.rand(4, 2, 1, 3, 3)
    print(i)
    m.eval()
    with torch.no_grad():
        o = m(i)
    print(o)
    print(i.shape, o.shape)


if __name__ == "__main__":
    main()