def compute_knn(distance_matrix: np.array,
k: int = 100) -> Tuple[np.array, np.array]:
k += 1 # k nearest neighbors + the element itself
k_i = distance_matrix.argpartition(k, axis=0)
k_d = np.take_along_axis(distance_matrix, k_i, axis=0)
sorted_indices = k_d.argsort(axis=0)
k_i_sorted = np.take_along_axis(k_i, sorted_indices, axis=0)[1:k]
k_d_sorted = np.take_along_axis(distance_matrix, k_i_sorted, axis=0)
knn_indices = np.transpose(k_i_sorted)
knn_distances = np.transpose(k_d_sorted)
return knn_distances, knn_indices
