Skip to content

Python API Reference

Generated from the docstrings in python/jasper/__init__.py.

Graph

Graph 

Graph(handle: int, config_id: str, data_type: DataType, distance: DistanceFunc, n_neighbors: int, dim: int, *, is_directional: bool = False, k_ranks: int | None = None, has_lsh: bool = False, has_pq: bool = False, prerotate: bool = False, prerotate_seed: int = 42)

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)

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.

  • directional_search –

    Run beam search scored by the cross-polytope LSH estimator.

  • pq_search –

    Run beam search scored by the Product-Quantization ADC estimator.

  • get_vector –

    Return the vector with the given stable id.

  • reserve_ids –

    Reserve count fresh 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:

  • n_tombstoned (int) –

    Number of soft-deleted vectors not yet consolidated away.

  • n_live (int) –

    Number of vectors still live (n_vectors - n_tombstoned).

n_tombstoned property 

n_tombstoned: int

Number of soft-deleted vectors not yet consolidated away.

n_live property 

n_live: int

Number of vectors still live (n_vectors - n_tombstoned).

load classmethod 

load(path: str, dim: int, n_neighbors: int = 32, data_type: DataType | str = FLOAT16, distance: DistanceFunc | str = L2, on_host: bool = False, k_ranks: int | None = None, prerotate: bool = False, prerotate_seed: int = 42) -> 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.

  • k_ranks (int | None, default: None ) –

    Pass the k_ranks used at build() time to load a directional graph. Its on-disk trailer (see Graph.build) is read automatically: whichever of directional_search()/pq_search() have their artifacts present become usable. Leave None to load a plain graph.

  • prerotate (bool, default: False ) –

    Must match what build() used for this graph — not stored in the file. Only relevant when k_ranks is given. FOOTGUN: build()'s default is prerotate = build_lsh or build_pq (usually True), but this default is False. If you built with prerotate on (the common case) and load() without passing prerotate=True, queries silently skip rotation against a rotated index — results are wrong with no error. Always pass the same prerotate/prerotate_seed used at build() time. (Not stored in the trailer yet — see save_directional_graph_to_file/ load_directional_graph_from_file in graph.cuh.)

  • prerotate_seed (int, default: 42 ) –

    Must match what build() used for this graph.

save 

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.

build classmethod 

build(vectors: Tensor, n_neighbors: int = 32, distance: DistanceFunc | str = L2, alpha: float = 1.2, workspace_budget: str = '10GB', on_host: bool = False, build_lsh: bool = False, build_pq: bool = False, k_ranks: int = 4, prerotate: bool | None = None, prerotate_seed: int = 42, lsh_samples: int = 32768, lsh_seed: int = 42, pq_train: int = 40000, pq_kmeans_iter: int = 12, pq_seed: int = 123) -> 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.

  • build_lsh (bool, default: False ) –

    Also populate cross-polytope LSH edges, enabling directional_search().

  • build_pq (bool, default: False ) –

    Also populate Product-Quantization edges + exact vector norms, enabling pq_search().

  • k_ranks (int, default: 4 ) –

    LSH rank count / PQ subquantizer count (4 or 16). Only used when build_lsh or build_pq is True.

  • prerotate (bool | None, default: None ) –

    Rotate vectors (and, transparently, queries at search time) by a random orthogonal matrix. Defaults to True iff build_lsh or build_pq is set, since LSH/PQ estimator quality relies on it — pass False explicitly to disable. NOTE: this choice (and prerotate_seed) is NOT persisted by save()/load() — see the ⚠ FOOTGUN note on Graph.load's prerotate arg.

  • prerotate_seed (int, default: 42 ) –

    Seed for the rotation matrix.

  • lsh_samples (int, default: 32768 ) –

    Edge samples used to calibrate LSH globals.

  • lsh_seed (int, default: 42 ) –

    Seed for LSH global sampling.

  • pq_train (int, default: 40000 ) –

    Residual samples used to train PQ codebooks.

  • pq_kmeans_iter (int, default: 12 ) –

    K-means iterations for PQ codebook training.

  • pq_seed (int, default: 123 ) –

    Seed for PQ codebook training.

Returns:

  • Graph –

    A Graph ready for search() (plain graphs), or additionally for

  • Graph –

    directional_search()/pq_search() when build_lsh/build_pq are set.

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]

directional_search(queries: Tensor, k: int = 10, beam_width: int = 64, limit: int | None = None, print_throughput: bool = False) -> tuple[Tensor, Tensor]

Run beam search scored by the cross-polytope LSH estimator.

Requires this graph to have been built with build_lsh=True (or loaded with LSH artifacts present — see Graph.load's k_ranks arg). Query rotation (if this graph was prerotated) happens transparently.

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.

  • limit (int | None, default: None ) –

    Defaults to 2x beam_width.

Returns: indices: int32 tensor [n_queries, k] distances: float32 tensor [n_queries, k]

pq_search(queries: Tensor, k: int = 10, beam_width: int = 64, limit: int | None = None, print_throughput: bool = False) -> tuple[Tensor, Tensor]

Run beam search scored by the Product-Quantization ADC estimator.

Requires this graph to have been built with build_pq=True (or loaded with PQ artifacts present — see Graph.load's k_ranks arg). Query rotation (if this graph was prerotated) happens transparently.

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.

  • limit (int | None, default: None ) –

    Defaults to 2x beam_width.

Returns: indices: int32 tensor [n_queries, k] distances: float32 tensor [n_queries, k]

get_vector 

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 

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 

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.

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 

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 

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 

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).

Enums

DistanceFunc

Bases: str, Enum

DataType

Bases: str, Enum

Vector I/O

read_bin 

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 

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 

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 

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]

get_recall 

get_recall(gt, result_indices, k, n_queries)

Utilities

parse_storage_size 

parse_storage_size(s: str) -> int