# MaxSim

> NumKong for Python

MaxSim is the late-interaction primitive used by systems such as [ColBERT](https://arxiv.org/abs/2004.12832).
It is not generic matrix multiplication.

```python
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")
```

`import numkong`

`from numkong import MaxSimPackedMatrix`

## MaxSimPackedMatrix

Opaque pre-packed matrix for MaxSim late-interaction scoring.

```python
class MaxSimPackedMatrix
```

## pack_size

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

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

## __repr__

Return a string representation.

```python
def __repr__() -> str
```

## vectors

Number of vectors.

```python
vectors: int
```

## depth

Number of dimensions per vector (depth).

```python
depth: int
```

## dtype

Data type of the packed vectors.

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

## nbytes

Size of the packed buffer in bytes.

```python
nbytes: int
```

## shape

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

```python
shape: tuple[int, int]
```

## maxsim_pack

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

```python
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.

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

## maxsim_packed

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

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

Returns the sum of per-query minimum angular distances.

- `queries`: Query vectors packed by maxsim_pack().
- `documents`: Document vectors packed by maxsim_pack().
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## maxsim

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

```python
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.

- `queries`: Query matrix of shape [query_count,depth].
- `documents`: Document matrix of shape [document_count,depth].
- `dtype`: Packing dtype, inferred from the inputs when None.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.