def init_rp_tree_map(index):
rows = chunk_rows(chunk_size, index, n_vertices)
return (
index,
sparse_init_rp_tree_map_jit(
rows,
leaf_array,
inds,
indptr,
data,
heap_updates[index],
dist,
),
)
def current_graph_map(index):
rows = chunk_rows(chunk_size, index, n_vertices)
return (
index,
sparse_current_graph_map_jit(
current_graph,
rows,
n_neighbors,
inds,
indptr,
data,
rng_state_threads[index],
seed_per_row=seed_per_row,
sparse_dist=dist,
),
)
def nn_descent_map(index):
rows = chunk_rows(chunk_size, index, n_vertices)
return (
index,
sparse_nn_descent_map_jit(
rows,
max_candidates,
inds,
indptr,
data,
new_candidate_neighbors,
old_candidate_neighbors,
heap_updates[index],
offset=0,
sparse_dist=dist,
),
)
def deheap_sort_map(index):
rows = chunk_rows(chunk_size, index, n_vertices)
return index, deheap_sort_map_jit(rows, current_graph)
