# MaxSim

> NumKong for Swift

MaxSim is the late-interaction primitive used by systems such as [ColBERT](https://arxiv.org/abs/2004.12832).
Each query is a small matrix of token vectors.
Each document is a small matrix of token vectors.
The score between a query and a document is the sum of maximum cosine similarities between each query token and any document token.
That is not a standard matrix multiply.

```swift
import NumKong

// 4 query tokens, each 16-dimensional
let queries = try Tensor<Float32>.full(rows: 4, columns: 16, value: 1.0)

// 8 document tokens, each 16-dimensional
let docs = try Tensor<Float32>.full(rows: 8, columns: 16, value: 1.0)

let queryPacked = try queries.maxSimPack()  // MaxSimPackedMatrix<Float32>
let docPacked   = try docs.maxSimPack()     // MaxSimPackedMatrix<Float32>

let score = try queryPacked.score(docPacked)  // Float64
assert(score.isFinite)
```

`MaxSimPackedMatrix` can also be constructed directly from a `MatrixView`:

```swift
let view = queries.view()
let packed = try MaxSimPackedMatrix<Float32>(packing: view)
```

Supported types and their output types:

| Input type | Score output |
| :--------- | :----------- |
| `Float32`  | `Float64`    |
| `BFloat16` | `Float32`    |
| `Float16`  | `Float32`    |

`Float16` support is unavailable on Intel macOS and Mac Catalyst, where Swift's `Float16` type does not exist; every other target, including Linux and Windows on x86-64, has it.

`import NumKong`

## NumKongMaxSimElement

Element type that supports late-interaction MaxSim scoring with packed representations.

```swift
protocol NumKongMaxSimElement
```

Conforming types: `[BFloat16](https://www.unum.cloud/docs/numkong/swift/scalar-types.md#BFloat16)`, [`Float`](https://www.unum.cloud/docs/numkong/swift/scalar-types.md#Float), [`Float16`](https://www.unum.cloud/docs/numkong/swift/scalar-types.md#Float16)

## MaxSimOutput

The type of the score that `MaxSimPackedMatrix.score(_:)` returns.

```swift
associatedtype MaxSimOutput
```

## MaxSimPackedMatrix

Owns a packed representation of multi-vector embeddings for MaxSim scoring.

```swift
final class MaxSimPackedMatrix<Element> where Element : NumKongMaxSimElement
```

Conforms to: `Sendable`

## vectors

```swift
let vectors: Int
```

## depth

```swift
let depth: Int
```

## byteCount

```swift
let byteCount: Int
```

## init

Packs a matrix view into the MaxSim-optimized layout.

```swift
convenience init(packing matrix: MatrixView<Element>) throws
```

## score

Computes the MaxSim score between this query and a document's packed matrix.

```swift
func score(_ document: MaxSimPackedMatrix<Element>) throws -> Element.MaxSimOutput
```

## shape

Reads the packed shape, __[vectors,depth]__, from the buffer's self-describing header.

```swift
var shape: (vectors: Int, depth: Int) { get throws }
```