Implement Barnes-Hut method

[alonfnt]

Currently, only the exact method is implemented, but realistically, one wants to avoid computing negligible terms whenever possible.

Most libraries rely on near neighbours to calculate the distances and joint probability.
Accelerating t-SNE using Tree-Based Algorithms

Note: Tree-based methods are annoying with JAX static shape requirements. Maybe lax.top_k is good enough here (although still O(n2)), or maybe we should rely on another library and not propagate gradients.