scib_metrics.nearest_neighbors.jax_approx_min_k

scib_metrics.nearest_neighbors.jax_approx_min_k(X, n_neighbors, recall_target=0.95, chunk_size=2048)

Run approximate nearest neighbor search using jax.

On TPU backends, this is approximate nearest neighbor search. On other backends, this is exact nearest neighbor search.

Parameters:

X ndarray

Data matrix.

n_neighbors int

Number of neighbors to search for.

recall_target float (default: 0.95)

Target recall for approximate nearest neighbor search.

chunk_size int (default: 2048)

Number of query points to search for at once.

Return type:

NeighborsResults