MaxSim | NumKong for Python

MaxSim is the late-interaction primitive used by systems such as ColBERT. It is not generic matrix multiplication.

import numpy as np
import numkong as nk

queries = np.random.randn(32, 128).astype(np.float32)
documents = np.random.randn(192, 128).astype(np.float32)

q = nk.maxsim_pack(queries, dtype="float32")
d = nk.maxsim_pack(documents, dtype="float32")
score = nk.maxsim_packed(q, d)

assert np.isfinite(score)
assert q.nbytes == nk.MaxSimPackedMatrix.pack_size(32, 128, dtype="float32")

API

ClassMaxSimPackedMatrixOpaque pre-packed matrix for MaxSim late-interaction scoring.7

__init__.pyi:589 ↗

class MaxSimPackedMatrix

pack_sizeReturn packed buffer size in bytes for given dimensions and dtype.

__init__.pyi:617 ↗

@classmethod
def pack_size(
    vectors: int,
    depth: int,
    /,
    dtype: _FloatTypeName | _MiniFloatType,
) -> int

__repr__Return a string representation.

__init__.pyi:622 ↗

def __repr__() -> str

vectorsNumber of vectors.

__init__.pyi:593 ↗

vectors: int

depthNumber of dimensions per vector (depth).

__init__.pyi:598 ↗

depth: int

dtypeData type of the packed vectors.

__init__.pyi:603 ↗

dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName

nbytesSize of the packed buffer in bytes.

__init__.pyi:608 ↗

nbytes: int

shapeDimensions (vectors, depth) read from the packed buffer header.

__init__.pyi:613 ↗

shape: tuple[int, int]

Functions

maxsim_packPack a matrix for MaxSim late-interaction scoring with maxsim_packed().

__init__.pyi:2277 ↗

def maxsim_pack(
    b: _BufferType,
    /,
    dtype: _FloatTypeName | _MiniFloatType | None = None,
    *,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> MaxSimPackedMatrix

Packing fuses no transposition, so transpose a depth-major matrix into a contiguous copy first.

Parameters

`b_BufferType`		Source matrix of shape [vectors,depth] with contiguous rows.
`dtype_FloatTypeName \| _MiniFloatType \| None`		Packing dtype, inferred from `b` when None.
`capabilitiesCapability \| int \| None`		Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
`streamint \| None`		A GPU stream pointer, left None on the CPU.

maxsim_packedCompute the MaxSim late-interaction score of two packed matrices.

__init__.pyi:2290 ↗

def maxsim_packed(
    queries: MaxSimPackedMatrix,
    documents: MaxSimPackedMatrix,
    /,
    *,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> float

Returns the sum of per-query minimum angular distances.

Parameters

`queriesMaxSimPackedMatrix`		Query vectors packed by maxsim_pack().
`documentsMaxSimPackedMatrix`		Document vectors packed by maxsim_pack().
`capabilitiesCapability \| int \| None`		Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
`streamint \| None`		A GPU stream pointer, left None on the CPU.

maxsimPack both matrices and compute their MaxSim late-interaction score in one call.

__init__.pyi:2383 ↗

def maxsim(
    queries: _BufferType,
    documents: _BufferType,
    /,
    dtype: _FloatTypeName | _MiniFloatType | None = None,
    *,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> float

Returns the sum of per-query minimum angular distances.

Parameters

`queries_BufferType`		Query matrix of shape [query_count,depth].
`documents_BufferType`		Document matrix of shape [document_count,depth].
`dtype_FloatTypeName \| _MiniFloatType \| None`		Packing dtype, inferred from the inputs when None.
`capabilitiesCapability \| int \| None`		Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
`streamint \| None`		A GPU stream pointer, left None on the CPU.