def generate_hash(n_points, d, b):
x = torch.rand(n_points, d)
a = torch.randn(b, d + 1)
a[:, d] = 0
h = torch.zeros(n_points, dtype=torch.int64)
compute_hashes(x, a, h)
return h
def cluster_queries(Q, query_lengths, C, I, B):
N, H, L, E = Q.shape
planes = Q.new_empty((B, E + 1))
normal_(planes)
planes[:, -1] = 0
hashes = compute_hashes(Q.view(N * H * L, E), planes).view(N, H, L)
# Cluster the hashes and return the cluster index per query
groups, counts = cluster(hashes,
query_lengths,
clusters=C,
iterations=I,
bits=B)
return groups, counts
def generate_hash(n_points, d, b, h):
torch.manual_seed(0)
x = torch.rand(n_points, d).cuda()
a = torch.randn(b, d + 1).cuda()
compute_hashes(x, a, h)
return h
