The dense distance entrypoints cover sqeuclidean, euclidean, and angular.
The first important difference from NumPy or SciPy is that the accumulator policy is not forced to match the storage dtype.
import numpy as np
import numkong as nk
a = np.random.randn(768).astype(np.float16)
b = np.random.randn(768).astype(np.float16)
sqeuclidean = nk.sqeuclidean(a, b)
euclidean = nk.euclidean(a, b)
angular = nk.angular(a, b)For float16, a naive same-dtype implementation is exactly the kind of path that loses precision or widens too late.
NumKong’s API makes the widening policy part of the kernel contract.
out=, dtype=, and out_dtype=Most distance and dot-product entrypoints accept out=, dtype=, and out_dtype= keyword arguments.
Passing them avoids dynamic memory allocations for temporary objects.
import numpy as np
import numkong as nk
queries = np.random.randn(100, 768).astype(np.float32)
database = np.random.randn(100, 768).astype(np.float32)
# Pre-allocated output with out=
out = nk.zeros((100,), dtype="float32")
nk.sqeuclidean(queries, database[:100], out=out) # writes in-place, returns None
# Explicit input dtype for raw byte buffers
raw = np.frombuffer(some_bytes, dtype=np.uint16)
nk.dot(raw, raw, dtype=nk.bfloat16) # reinterpret uint16 as bf16
# Output dtype override
nk.euclidean(queries[0], database[0], out_dtype="float32") # accumulate in f64, downcast resultWhen out= is provided, the function writes results in-place and returns None.
The out array must be pre-allocated with the correct shape and a supported dtype.
For custom float types (bfloat16, float16, float8_e4m3, float8_e5m2, float6_e2m3, float6_e3m2), type objects are preferred over strings — they are faster to dispatch and provide IDE autocomplete:
nk.dot(a, b, dtype=nk.bfloat16) # works faster
nk.dot(a, b, dtype="bfloat16") # works a bit slowerpointer_to_*× 3Return the address of the euclidean kernel for dtype, as an integer.def pointer_to_euclidean(
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType,
/,
) -> int| Symbol |
|---|
pointer_to_euclidean |
pointer_to_sqeuclidean |
pointer_to_angular |
sqeuclideanCompute the squared Euclidean distance between two vectors.def sqeuclidean(
a: _BufferType,
b: _BufferType,
/,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType | None = None,
*,
out: _BufferType | None = None,
out_dtype: _FloatTypeName | _MiniFloatType | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> float | Tensor | NoneSimilar to scipy.spatial.distance.sqeuclidean. Two vectors give a scalar. Matrices pair row i
of a with row i of b, a single row pairing with every row of the other, into a new Tensor of
one result per pair, or into out, returning None.
a_BufferType | First real or integer vector or matrix. | |
|---|---|---|
b_BufferType | Second real or integer vector or matrix, of the dtype of a. | |
dtypeLiteral['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType | None | Input dtype to read both buffers as, overriding the one they report. | |
out_BufferType | None | Output vector for the results. | |
out_dtype_FloatTypeName | _MiniFloatType | None | Result dtype, by default that of out, else float64. | |
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. |
angularCompute the angular distance, also known as cosine distance, between two vectors.def angular(
a: _BufferType,
b: _BufferType,
/,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType | None = None,
*,
out: _BufferType | None = None,
out_dtype: _FloatTypeName | _MiniFloatType | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> float | Tensor | NoneSimilar to scipy.spatial.distance.cosine. Two vectors give a scalar. Matrices pair row i of
a with row i of b, a single row pairing with every row of the other, into a new Tensor of
one result per pair, or into out, returning None.
a_BufferType | First real or integer vector or matrix. | |
|---|---|---|
b_BufferType | Second real or integer vector or matrix, of the dtype of a. | |
dtypeLiteral['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType | None | Input dtype to read both buffers as, overriding the one they report. | |
out_BufferType | None | Output vector for the results. | |
out_dtype_FloatTypeName | _MiniFloatType | None | Result dtype, by default that of out, else float64. | |
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. |
euclideanCompute the Euclidean distance between two vectors.def euclidean(
a: _BufferType,
b: _BufferType,
/,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType | None = None,
*,
out: _BufferType | None = None,
out_dtype: _FloatTypeName | _MiniFloatType | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> float | Tensor | NoneSimilar to scipy.spatial.distance.euclidean. Two vectors give a scalar. Matrices pair row i of
a with row i of b, a single row pairing with every row of the other, into a new Tensor of
one result per pair, or into out, returning None.
a_BufferType | First real or integer vector or matrix. | |
|---|---|---|
b_BufferType | Second real or integer vector or matrix, of the dtype of a. | |
dtypeLiteral['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType | None | Input dtype to read both buffers as, overriding the one they report. | |
out_BufferType | None | Output vector for the results. | |
out_dtype_FloatTypeName | _MiniFloatType | None | Result dtype, by default that of out, else float64. | |
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. |