I have a dataset of a thousand 128 dimensional features in the shape of e.g. (1000,128).

I want to find the sorted nearest neighbors of a 128 dimensional feature in the shape of (128,1).

The distance in calculated via a Matrix Multiplication between dataset (1000,128) and feature (128,1) which would give an array of similarities in the shape of (1000,1) :

## DATASET (1000,128) x FEATURE (128,1) = SIMILARITIES (1000,1)

This is done via:

```
# features.shape=(1000,128) ; feature.shape=(128,1) ; similarities.shape=(1000,1)
similarities = features.dot(feature)
```

After calculating the distance (similarities), I'm finding the nearest neighbors using the code below:

```
# The n Nearest Neighbors Indexes (But Not Sorted)
nearest_neighbours_indexes_unsorted = np.argpartition(similarities, kth=-n)[-n:]
# The n Nearest Neighbors (But Not Sorted)
nearest_neighbours_similarities_unsorted = similarities[nearest_neighbours_indexes_unsorted]
# The Indexes of n Nearest Neighbors Sorted
nearest_neighbours_indexes_sorted = np.flip(nearest_neighbours_indexes_unsorted[np.argsort(nearest_neighbours_similarities_unsorted)], axis=0)
```

This code works very fast for millions of data (I'm interested if someone has a tip to make it faster) But I want to be able to find the nearest neighbors of more than one feature in one go:

## DATASET (1000,128) x FEATURE (128,n) = SIMILARITIES (1000,n)

One way is to calculate the above code for each feature in a loop (which is slow) and the other way is to change the code to accommodate for multidimensional indexing and here's where I'm stuck: I don't know how to write the above code for features in the shape of (128,n) and not (128,1).

## Helper functions to get largest, smallest n-indices, elements along an axis

Here's a helper function to select top

`n-largest`

indices along a generic axis from a generic ndarray making use of`np.argpartition`

and`np.take_along_axis`

-Extending this to get n-smallest indices -

And extending these to select the largest or smallest

`n`

elements themselves, it would be with a simple usage of`np.take_along_axis`

as listed next -Sample runsTake largest

`n`

indices, elements along axis -Take smallest

`n`

indices, elements along axis -## Solving our case here

For our case, let's assume the input is

`similarities`

and is of shape`(1000,128)`

representing 1000 data points and 128 features and that we want to look for largest say`n=10`

features for each of those data points, then it would be -The final indices/values array would be of shape

`(1000, n)`

.Sample run with the given dataset shape -

If instead you were looking to get

`n`

largest data-points for each of those features, that is the final indices/values array would be of shape`(n, 128)`

, then use`axis=0`

.