import torch

n = 10**5
k = 59
batch_size = 10**3
device = 'cuda:0' # 'cpu'
nn_number = 3

items = np.random.rand(n, k).astype(np.float32)
items = items / np.sqrt(np.sum(items**2, axis=1, keepdims=True))

def three_closest(X, indices):
    sample_n = indices.shape[0]
    index = torch.zeros((sample_n, nn_number), dtype=torch.int64, device=device)
    distances = X[indices, :] @ X.T
    ind_range = torch.arange(sample_n, device=device)
    distances[ind_range, indices] = 0
    for j in range(nn_number):
        max_index = torch.argmax(distances, dim=1)
        index[:, j] = max_index
        distances[ind_range, max_index] = 0
    return index

def distance_func(X):
    X = torch.from_numpy(X).to(device=device)
    all_items = torch.zeros((n, nn_number), dtype=torch.int64, device=device)
    for indices in torch.split(torch.arange(n, device=device), batch_size):        
        all_items[indices] = three_closest(X, indices)
    return all_items.cpu().numpy()

nearest_torch = distance_func(items)

import numpy as np
import faiss


n = 10**5
k = 59
batch_size = 18

items = np.random.rand(n, k).astype(np.float32)
items = items / np.sqrt(np.sum(items**2, axis=1, keepdims=True))
bounds = np.linspace(0, n, n // batch_size + 1, dtype=np.int64)

index = faiss.IndexFlatIP(59)
index.add(items)

distance = np.zeros((n, 4), dtype=np.float32)
index_val = np.zeros((n, 4), dtype=np.int64)
for lower, upper in zip(bounds[:-1], bounds[1:]):
    distance[lower:upper], index_val[lower:upper] = index.search(items[lower:upper], 4)
assert np.sum(index_val[:, 0] != np.arange(n)) == 0
distance = distance[:, 1:]
index_val = index_val[:, 1:]