Python API Reference¶
Generated from the docstrings in python/jasper/__init__.py.
Graph¶
Graph(handle: int, config_id: str, data_type: DataType, distance: DistanceFunc, n_neighbors: int, dim: int)
A GPU-resident graph index for nearest neighbor search.
Usage
Load from file¶
g = jasper.Graph.load("sift1m.graph", dim=128, n_neighbors=32)
Build from vectors¶
vectors = torch.randn(100000, 128, device="cuda") g = jasper.Graph.build(vectors, n_neighbors=32)
Search¶
indices, distances = g.search(queries, k=10)
Search over RaBitQ codes, with an exact rerank of the final beam¶
g.build_rabitq(bits=4) indices, distances = g.search_rabitq(queries, k=10)
Methods:
-
load–Load a graph from a binary file into GPU memory.
-
save–Save this graph to a binary file.
-
build–Construct a graph index from vectors on GPU.
-
search–Run beam search on this graph.
-
build_rabitq–Quantize the graph's vectors with RaBitQ, enabling
search_rabitq. -
free_rabitq–Free the RaBitQ codes built by
build_rabitq; the graph is kept. -
search_rabitq–Run beam search scored by RaBitQ distance estimates, then rerank the
-
get_vector–Return the vector with the given stable id.
-
reserve_ids–Reserve
countfresh stable ids, returning them as an int32 CPU -
append–Append a batch of vectors to the live graph, wiring their edges into the
-
mark_deleted–Soft-delete a batch of vectors by id.
-
consolidate–Repair graph edges that route through deleted vertices and clear all
-
compact–Reclaim space by compacting live vectors into internal slots
Attributes:
-
has_rabitq(bool) –Whether RaBitQ codes are built and current (see
build_rabitq). -
n_tombstoned(int) –Number of soft-deleted vectors not yet consolidated away.
-
n_live(int) –Number of vectors still live (n_vectors - n_tombstoned).
property
¶
has_rabitq: bool
Whether RaBitQ codes are built and current (see build_rabitq).
classmethod
¶
load(path: str, dim: int, n_neighbors: int = 32, data_type: DataType | str = FLOAT16, distance: DistanceFunc | str = L2, on_host: bool = False) -> Graph
Load a graph from a binary file into GPU memory.
Parameters:
-
path(str) –Path to the graph binary file.
-
dim(int) –Dimensionality of the vectors.
-
n_neighbors(int, default:32) –Max neighbors per node (must match file).
-
data_type(DataType | str, default:FLOAT16) –Vector data type: "f16".
-
distance(DistanceFunc | str, default:L2) –Distance function: "l2" or "ip".
-
on_host(bool, default:False) –Load the graph on host memory.
save(path: str) -> None
Save this graph to a binary file.
The file can later be reloaded with Graph.load(), using the
same dim, n_neighbors, data_type, and distance
that were used when the graph was built or originally loaded.
Parameters:
-
path(str) –Destination file path.
classmethod
¶
build(vectors: Tensor, n_neighbors: int = 32, distance: DistanceFunc | str = L2, alpha: float = 1.2, workspace_budget: str = '10GB', on_host: bool = False) -> Graph
Construct a graph index from vectors on GPU.
Parameters:
-
vectors(Tensor) –CUDA tensor of shape [n_vectors, dim].
-
n_neighbors(int, default:32) –Max neighbors per node (R).
-
distance(DistanceFunc | str, default:L2) –Distance function: "l2" or "ip".
-
alpha(float, default:1.2) –Pruning factor (1.0 = strict, >1.0 = long hops).
-
workspace_budget(str, default:'10GB') –GPU memory budget (default to 10GB).
-
on_host(bool, default:False) –Construct the graph on host memory.
Returns:
-
Graph–A Graph ready for search.
search(queries: Tensor, k: int = 10, beam_width: int = 64, print_throughput: bool = False) -> tuple[Tensor, Tensor]
Run beam search on this graph.
Parameters:
-
queries(Tensor) –CUDA tensor of shape [n_queries, dim].
-
k(int, default:10) –Number of nearest neighbors to return.
-
beam_width(int, default:64) –Search beam width.
Returns: indices: int32 tensor [n_queries, k] of stable ids distances: float32 tensor [n_queries, k]
build_rabitq(bits: int = 4, seed: int = 0) -> None
Quantize the graph's vectors with RaBitQ, enabling search_rabitq.
Each vector is centered on the dataset centroid, randomly rotated, and
stored as a bits-per-dimension code plus two float factors (72
bytes per 128-dim vector at 4 bits, versus 256 bytes as float16). The
graph itself is unchanged; the codes live next to it on the GPU.
Calling this again replaces the previous codes.
The codes snapshot the graph's current slots. append, and a
compact that reclaims slots, discard them (has_rabitq becomes
False), so call build_rabitq again afterwards. mark_deleted and
consolidate keep them valid.
Parameters:
-
bits(int, default:4) –Bits per dimension: 1, 2, 4 or 8. More bits give tighter distance estimates at a higher memory and bandwidth cost.
-
seed(int, default:0) –Seed for the random rotation.
search_rabitq(queries: Tensor, k: int = 10, beam_width: int = 64, print_throughput: bool = False) -> tuple[Tensor, Tensor]
Run beam search scored by RaBitQ distance estimates, then rerank the final beam with exact distances.
The graph walk reads the compact RaBitQ codes instead of the full
vectors; only the beam_width surviving candidates are rescored
against the original vectors, so the returned distances are exact and
directly comparable to search. Requires build_rabitq() first.
Parameters:
-
queries(Tensor) –CUDA tensor of shape [n_queries, dim].
-
k(int, default:10) –Number of nearest neighbors to return.
-
beam_width(int, default:64) –Search beam width (also the number of reranked candidates).
Returns: indices: int32 tensor [n_queries, k] of stable ids distances: float32 tensor [n_queries, k] (exact)
get_vector(stable_id: int) -> Tensor
Return the vector with the given stable id.
Stable ids are assigned in monotonic order as vectors are added and are
unchanged by consolidate/compact — they are not bounded by
n_vectors. Raises if the id is not present (deleted or never assigned).
Parameters:
-
stable_id(int) –The vector's stable id (as returned by
search).
Returns:
-
Tensor–A 1-D CUDA tensor of shape [dim] with the graph's data type.
reserve_ids(count: int) -> Tensor
Reserve count fresh stable ids, returning them as an int32 CPU
tensor of shape [count]. Advances the monotonic id counter; ids are
never reused. This is the id half of a live append — the caller writes
the corresponding vectors into the graph and registers each (id, slot).
Returns:
-
Tensor–int32 tensor [count] of newly assigned stable ids.
append(vectors: Tensor, alpha: float = 1.2) -> Tensor
Append a batch of vectors to the live graph, wiring their edges into the
existing graph (beam-search + robust-prune, same as construction) and
assigning each a fresh monotonic stable id. Discards any RaBitQ codes;
call build_rabitq again to keep using search_rabitq.
Parameters:
-
vectors(Tensor) –CUDA tensor [n, dim] of the graph's data type (float16).
-
alpha(float, default:1.2) –Robust-pruning factor for the new vectors' edges.
Returns:
-
Tensor–int32 tensor [n] of the assigned stable ids, in input order.
mark_deleted(ids: Tensor) -> None
Soft-delete a batch of vectors by id.
Deleted vectors are immediately excluded from search results.
Their graph edges are repaired lazily by consolidate and their
slots reclaimed by compact. Out-of-range ids are ignored.
Parameters:
-
ids(Tensor) –1-D integer tensor of vector ids to delete.
consolidate(alpha: float = 1.2) -> None
Repair graph edges that route through deleted vertices and clear all
tombstones. After this call n_tombstoned is 0; ids are unchanged.
Parameters:
-
alpha(float, default:1.2) –Robust-pruning factor used when re-selecting edges.
compact() -> None
Reclaim space by compacting live vectors into internal slots
[0, n_live). Consolidates first if there are pending deletions.
Stable ids are preserved: a vector keeps the same id across compact
(only internal slots are renumbered, transparently via the id map).
If slots were reclaimed, RaBitQ codes are discarded; call
build_rabitq again to keep using search_rabitq.
Enums¶
Bases: str, Enum
Bases: str, Enum
Vector I/O¶
read_bin(path: str, dtype: str = 'f32', max_vectors: int = 0) -> Tensor
Read a [n, dim] binary file (f32 or u8 on disk) and return a pinned torch.float16 tensor.
Ground truth¶
read_groundtruth(path: str, k: int = 10) -> tuple[Tensor, Tensor]
Read ground truth from a binary file.
Format: [n_queries: uint32][gt_k: uint32][ids: n_queries * gt_k * uint32][distances: n_queries * gt_k * float32]
Parameters:
-
path(str) –Path to the ground truth .bin file.
-
k(int, default:10) –Number of neighbors to return (must be <= gt_k in file).
Returns:
-
indices(Tensor) –int32 tensor [n_queries, k]
-
distances(Tensor) –float32 tensor [n_queries, k]
generate_groundtruth(vectors: Tensor, queries: Tensor, k: int = 100, distance: str = 'l2', query_batch_size: int = 1024, vector_batch_size: int = 100000, device: str = 'cuda') -> tuple[Tensor, Tensor]
Brute-force exact k-NN on GPU, streaming both vectors and queries in batches to stay within device memory.
Parameters:
-
vectors(Tensor) –[n, dim] CPU tensor (float32 or float16) — the database.
-
queries(Tensor) –[nq, dim] CPU tensor (float32 or float16) — the queries.
-
k(int, default:100) –Number of nearest neighbors.
-
distance(str, default:'l2') –"l2" or "ip" (inner product).
-
query_batch_size(int, default:1024) –Queries transferred to device per batch.
-
vector_batch_size(int, default:100000) –Vectors transferred to device per batch.
-
device(str, default:'cuda') –Target device.
Returns:
-
indices(Tensor) –int32 [nq, k] (CPU)
-
distances(Tensor) –float32 [nq, k] (CPU)
save_groundtruth(path: str, indices: Tensor, distances: Tensor)
Write ground truth to the binary format expected by jasper.read_groundtruth.
Format: [n_queries: u32][k: u32][ids: u32 * nq * k][dists: f32 * nq * k]