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)

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.

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

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

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]

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