jaxknn¤

Exact k-nearest neighbours in JAX: a kd-tree that traces under jit, vmap and grad.

pip install jaxknn

import jax
import jax.numpy as jnp
import jaxknn

points = jax.random.normal(jax.random.key(0), (1000, 3))

# Each point's 5 nearest *other* points (self excluded by index).
idx, dist = jaxknn.knn(points, 5)
assert idx.shape == dist.shape == (1000, 5)

# The 3 nearest points to each query.
queries = jnp.zeros((2, 3))
idx, dist = jaxknn.knn(points, 3, queries=queries)
assert idx.shape == (2, 3)

Both return (indices, distances): - distances are Euclidean, with each row sorted ascending; - equidistant neighbours go to the lower index; - a missing neighbour (fewer than k candidates) is index len(points) with distance inf.

Where to go next¤

• Use it under jit, vmap and grad: the reason jaxknn exists.
• Find the nearest point with another label: for example, bridging the components of a graph.
• Tune it for large k or high dimension.
• How it works and when to use SciPy instead.
• Reference: jaxknn (the safe entry points) and jaxknn.tree (the raw tree).