Function knn

experiments/python/utils.py:123–126

(X, q, k, dist_func=dists_sq)

def knn(X, q, k, dist_func=dists_sq):
    dists = dist_func(X, q)
    idxs = top_k_idxs(dists, k)
    return idxs, dists[idxs]


@_memory.cache