def test_simple_product(self):
    N = 2
    H = 2
    L = 1000
    E = 32
    S = 1000
    k = 32
    C = 50
    I = 5
    B = 16
    Q = torch.randn(N, H, L, E).to(self.device)
    K = torch.randn(N, H, S, E).to(self.device)
    lengths = torch.full((N,), L).int().to(self.device)
    (groups, counts) = cluster_queries(Q, lengths, C, I, B)
    Q_grouped = aggregate(Q, groups, (1 / counts.float()))
    QK = torch.einsum('nhle,nhse->nhls', Q_grouped, K)
    (_, topk) = torch.topk(QK, k, dim=(- 1))
    topk = topk.contiguous()
    products = clustered_sparse_dot_product(Q, K, topk, groups, counts, lengths)
    topk_broadcast = broadcast(topk.float(), groups, torch.ones_like(counts, dtype=torch.float32), torch.zeros((N, H, L, k), device=Q.device))
    all_products = torch.einsum('nhle,nhse->nhls', Q, K)
    products_2 = all_products[(torch.arange(N).view(N, 1, 1, 1), torch.arange(H).view(1, H, 1, 1), torch.arange(L).view(1, 1, L, 1), topk_broadcast.long())]
    self.assertLess(torch.max(torch.abs((products_2 - products))), 0.0001)
def test_simple_product(self):
    N = 2
    H = 2
    L = 1000
    E = 32
    S = 1000
    k = 32
    C = 50
    I = 5
    B = 16
    Q = torch.randn(N, H, L, E).to(self.device)
    K = torch.randn(N, H, S, E).to(self.device)
    lengths = torch.full((N,), L).int().to(self.device)
    (groups, counts) = cluster_queries(Q, lengths, C, I, B)
    Q_grouped = aggregate(Q, groups, (1 / counts.float()))
    QK = torch.einsum('nhle,nhse->nhls', Q_grouped, K)
    (_, topk) = torch.topk(QK, k, dim=(- 1))
    topk = topk.contiguous()
    products = clustered_sparse_dot_product(Q, K, topk, groups, counts, lengths)
    topk_broadcast = broadcast(topk.float(), groups, torch.ones_like(counts, dtype=torch.float32), torch.zeros((N, H, L, k), device=Q.device))
    all_products = torch.einsum('nhle,nhse->nhls', Q, K)
    products_2 = all_products[(torch.arange(N).view(N, 1, 1, 1), torch.arange(H).view(1, H, 1, 1), torch.arange(L).view(1, 1, L, 1), topk_broadcast.long())]
    self.assertLess(torch.max(torch.abs((products_2 - products))), 0.0001)