jz-tree

jz-tree offers a framwork for GPU-friendly implementations of tree algorithms in jax (with a CUDA backend). Currently, nearest neighbour search and friends-of-friends are implemented, but more may come in the future! As far as I know, jz-tree offers the fastest GPU implementation of these two algorithms at the time of writing. For example, this is a benchmark of a 30-nearest neighbour search in three dimensions, run on a single NVIDIA-A100 GPU:

The code is open source and can easily be installed with:

pip install jztree[cuda12]   # or jztree[cuda13]

For additional details, also consider:

• The official documentation
• The code repository
• The arXiv preprint article: arXiv:2604.05885

Other codes

More codes are available on my GitHub page