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