jz-tree documentation

Welcome to the documentation of 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 we 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:

Contents:

• Installation
  • Via pip
  • Build from sources
  • Hello World
• Quickstart
  • Nearest neighbour search
  • Friends-of-friends groups
  • Friends-of-friends catalogues
  • Error handling
• Advanced Guide: Multi-GPU
  • shard_map in JAX
  • Particle-data and padding
  • Distributed kNN
  • Distributed friends-of-friends
  • Multi-host execution and performance
• Developer Guide
  • Code structure
  • Python modules
  • CUDA kernels and automatic FFI generation
  • Unit tests and benchmarks
  • Recommendations
  • Building wheels
• API Reference
  • jztree.config
  • jztree.data
  • jztree.tree
  • jztree.knn
  • jztree.fof
• Attribution