jaxkd 0.0.1

Minimal JAX implementation of k-nearest neighbors using a k-d tree.

Minimal JAX implementation of k-nearest neighbors using a k-d tree!

This is essentially just a translation of two GPU-friendly tree algorithms [1, 2] into XLA primitives. It is convenient and lightweight, but the original CUDA implementation may be a better choice depending on the application.

The build_tree and query_neighbors operations are compatible with JIT and automatic differentiation. They are reasonably fast when vectorized on GPU, but will be much slower than scipy.spatial.KDTree on CPU. The main advantage is to avoid the complexity of using non-JAX libraries and potentially leaving JIT and the GPU when a scalable nearest neighbor search is needed as part of a larger JAX program.

Usage:

import jax
import jaxkd as jk
kp, kq = jax.random.split(jax.random.key(83))

points = jax.random.normal(kp, shape=(100_000, 3))
queries = jax.random.normal(kq, shape=(10_000, 3))
tree = jk.build_tree(points)
neighbors, distances = jk.query_neighbors(tree, queries, 10)

Notes:

• The tree structure is stored with a tuple of arrays: points in the original order (not copied), indices to put the points in tree order, and split_dims which (if the tree is built with optimized=True) specify the splitting dimension independently for each node. If needed, the memory overhead could potentially be reduced by sorting points in-place.
• The query_neighbors function is intended for relatively small values of k and does not use a max heap for simplicity. If k is large enough that will become worthwhile.
• The demo.ipynb notebook in the source repository has some additional examples, including gradient-based optimization using neighbors.