Create a numba accelerated version of nearest neighbor descent
    specialised for the given distance metric and metric arguments on sparse
    matrix data provided in CSR ind, indptr and data format. Numba
    doesn't support higher order functions directly, but we can instead JIT
    compile the version of NN-descent for any given metric.

    Parameters
    ----------
    sparse_dist: function
        A numba JITd distance function which, given four arrays (two sets of
        indices and data) computes a dissimilarity between them.

    dist_args: tuple
        Any extra arguments that need to be passed to the distance function
        beyond the two arrays to be compared.

    Returns
    -------
    A numba JITd function for nearest neighbor descent computation that is
    specialised to the given metric.

def make_sparse_nn_descent(sparse_dist, dist_args):
    """Create a numba accelerated version of nearest neighbor descent
    specialised for the given distance metric and metric arguments on sparse
    matrix data provided in CSR ind, indptr and data format. Numba
    doesn't support higher order functions directly, but we can instead JIT
    compile the version of NN-descent for any given metric.

    Parameters
    ----------
    sparse_dist: function
        A numba JITd distance function which, given four arrays (two sets of
        indices and data) computes a dissimilarity between them.

    dist_args: tuple
        Any extra arguments that need to be passed to the distance function
        beyond the two arrays to be compared.

    Returns
    -------
    A numba JITd function for nearest neighbor descent computation that is
    specialised to the given metric.
    """

    @numba.njit(parallel=True)
    def nn_descent(
        inds,
        indptr,
        data,
        n_vertices,
        n_neighbors,
        rng_state,
        max_candidates=50,
        n_iters=10,
        delta=0.001,
        rho=0.5,
        rp_tree_init=True,
        leaf_array=None,
        verbose=False,
    ):

        current_graph = make_heap(n_vertices, n_neighbors)
        for i in range(n_vertices):
            indices = rejection_sample(n_neighbors, n_vertices, rng_state)
            for j in range(indices.shape[0]):

                from_inds = inds[indptr[i] : indptr[i + 1]]
                from_data = data[indptr[i] : indptr[i + 1]]

                to_inds = inds[indptr[indices[j]] : indptr[indices[j] + 1]]
                to_data = data[indptr[indices[j]] : indptr[indices[j] + 1]]

                d = sparse_dist(from_inds, from_data, to_inds, to_data, *dist_args)

                heap_push(current_graph, i, d, indices[j], 1)
                heap_push(current_graph, indices[j], d, i, 1)

        if rp_tree_init:
            for n in range(leaf_array.shape[0]):
                for i in range(leaf_array.shape[1]):
                    if leaf_array[n, i] < 0:
                        break
                    for j in range(i + 1, leaf_array.shape[1]):
                        if leaf_array[n, j] < 0:
                            break

                        from_inds = inds[
                            indptr[leaf_array[n, i]] : indptr[leaf_array[n, i] + 1]
                        ]
                        from_data = data[
                            indptr[leaf_array[n, i]] : indptr[leaf_array[n, i] + 1]
                        ]

                        to_inds = inds[
                            indptr[leaf_array[n, j]] : indptr[leaf_array[n, j] + 1]
                        ]
                        to_data = data[
                            indptr[leaf_array[n, j]] : indptr[leaf_array[n, j] + 1]
                        ]

                        d = sparse_dist(
                            from_inds, from_data, to_inds, to_data, *dist_args
                        )

                        heap_push(
                            current_graph, leaf_array[n, i], d, leaf_array[n, j], 1
                        )
                        heap_push(
                            current_graph, leaf_array[n, j], d, leaf_array[n, i], 1
                        )

        for n in range(n_iters):
            if verbose:
                print("\t", n, " / ", n_iters)

            candidate_neighbors = build_candidates(
                current_graph, n_vertices, n_neighbors, max_candidates, rng_state
            )

            c = 0
            for i in range(n_vertices):
                for j in range(max_candidates):
                    p = int(candidate_neighbors[0, i, j])
                    if p < 0 or tau_rand(rng_state) < rho:
                        continue
                    for k in range(max_candidates):
                        q = int(candidate_neighbors[0, i, k])
                        if (
                            q < 0
                            or not candidate_neighbors[2, i, j]
                            and not candidate_neighbors[2, i, k]
                        ):
                            continue

                        from_inds = inds[indptr[p] : indptr[p + 1]]
                        from_data = data[indptr[p] : indptr[p + 1]]

                        to_inds = inds[indptr[q] : indptr[q + 1]]
                        to_data = data[indptr[q] : indptr[q + 1]]

                        d = sparse_dist(
                            from_inds, from_data, to_inds, to_data, *dist_args
                        )

                        c += heap_push(current_graph, p, d, q, 1)
                        c += heap_push(current_graph, q, d, p, 1)

            if c <= delta * n_neighbors * n_vertices:
                break

        return deheap_sort(current_graph)

    return nn_descent