# NumKong for Python > Mixed-precision math for every CPU and GPU, from Python It combines NumPy-friendly buffers with native mixed-precision kernels, zero-copy tensor views, packed and symmetric matrix operations, sparse helpers, many-to-many shape distances, and MaxSim. The API feels NumPy-shaped with familiar scalar, batched, and all-pairs entrypoints, while `Tensor` keeps shape, dtype, and strides visible through a memoryview-backed container. Low-precision dtypes (BFloat16, Float8, Float6, packed bits) flow through the same API, and dense, packed, and symmetric kernels release the GIL around native work. ## Installation From PyPI: ```sh python -m pip install numkong ``` From a local checkout: ```sh python -m pip install . ``` Quick runtime check: ```sh python -c "import numkong as nk; print(repr(nk.cpu_capabilities_enabled()))" ``` ### Wheel Compatibility and Building from Source Pre-built wheels are available on PyPI for Linux (x86_64, aarch64, riscv64, plus i686, ppc64le, s390x), macOS (x86_64, arm64), and Windows (AMD64, ARM64). Python 3.12 through 3.15 is supported, including the free-threaded 3.14t and 3.15t. Every wheel compiles each capability the toolchain supports and picks among them at runtime, so a single wheel covers all CPU generations on a given architecture. When building from source, the compiler requirements depend on the platform. On macOS x86 only AVX2 is available; on macOS ARM NEON is always present, but SME requires Apple M4+ with Xcode 16+ (AppleClang 16+). RISC-V builds require Clang and LLD because GCC lacks `zvfh`, `zvfbfwma`, and `zvbb` support. On Windows, MSVC 19.44+ (Visual Studio 2022 17.14+) is recommended for full AVX-512 with FP16/BF16/VNNI. Build parallelism is controlled by `CMAKE_BUILD_PARALLEL_LEVEL`, which should be lowered in memory-constrained containers. There is no OpenMP dependency. Python-side parallelism uses the `threads=` argument on the GIL-free kernels, or `concurrent.futures` around them. ```sh CMAKE_BUILD_PARALLEL_LEVEL=2 pip install . --no-build-isolation ``` ## Quickstart ```python import numpy as np import numkong as nk a, b = np.random.randn(1536).astype(np.float32), np.random.randn(1536).astype(np.float32) dot = nk.dot(a, b) # widened accumulation, not same-dtype print(dot) ``` ## Highlights ### Ecosystem Comparison | Feature | NumKong | [NumPy][numpy]/[SciPy][scipy] | [PyTorch][pytorch] | | :------------------------------ | :--------------------------------------------------------------------------------------------------------------------------- | :--------------------------------------------------------------------- | :-------------------------------------------------------------------------- | | Operation families | dots, distances, binary, probability, geospatial, curved, shapes, sparse, MaxSim, elementwise, reductions, cast, trig | dots, distances, elementwise, reductions, some probability via `cdist` | dots, distances, elementwise, reductions | | Precision | BFloat16 through sub-byte — Float8, Float6, Int4, packed bits; automatic widening; Kahan summation; 0 ULP in Float32/Float64 | Float16, partial BFloat16; no auto-widening; standard accuracy | Float16, BFloat16, partial Float8; explicit AMP required; standard accuracy | | Runtime SIMD dispatch | picks the best ISA per call, from a capability mask, at runtime on x86, ARM, RISC-V | runtime, among a few compiled targets per loop | CPU: compile-time; CUDA: runtime | | Packed matrix, GEMM-like | pack once, reuse across query batches | `np.dot`/`@` — no persistent packing | `torch.mm` — no persistent distance-oriented packing | | Symmetric kernels, SYRK-like | skip duplicate pairs, up to 2x speedup for self-distance | `pdist` computes one triangle; `cdist` recomputes both | `X @ X.T` recomputes both triangles | | Output parameter `out=` | Yes — all major entrypoints | Yes — most `ufunc`s and functions; SciPy: some functions only | Yes for `torch.mm`, `torch.matmul`; No for `torch.cdist` | | Fast CPython calling convention | Yes — direct `METH_FASTCALL` | Yes — `vectorcall` in 2.0+ | No — tensor dispatch overhead | | GIL release | batched, packed, and symmetric kernels | some ops only | most ops | [numpy]: https://github.com/numpy/numpy [scipy]: https://github.com/scipy/scipy [pytorch]: https://github.com/pytorch/pytorch ## Scalar Types NumKong exposes two different low-precision stories in Python. It exposes Python scalar objects for a few formats. And it exposes tensor dtypes for the broader buffer-oriented path. The six scalar types have stable payload sizes even though Python object headers are not: | Type | Bits | Bytes | Range | Inf | NaN | | :--------------- | :----- | ----: | --------: | :-: | :-: | | `nk.float16` | 1+5+10 | 2 | ±65504 | yes | yes | | `nk.bfloat16` | 1+8+7 | 2 | ±3.4×10³⁸ | yes | yes | | `nk.float8_e4m3` | 1+4+3 | 1 | ±448 | no | yes | | `nk.float8_e5m2` | 1+5+2 | 1 | ±57344 | yes | yes | | `nk.float6_e2m3` | 1+2+3 | 1 | ±7.5 | no | no | | `nk.float6_e3m2` | 1+3+2 | 1 | ±28 | no | no | The Bits column shows sign + exponent + mantissa bit counts. The Bytes column is the stable payload size; `float8_*` and `float6_*` both store 1 byte because the sub-byte formats are padded to byte alignment. The full object footprint is interpreter-dependent. Use `sys.getsizeof(nk.float16(1.0))` if you need the heap footprint of the Python wrapper object itself. Use `Tensor.itemsize` and `Tensor.nbytes` for the stable payload sizes of array storage. `ml_dtypes` matters here because NumKong explicitly interoperates with the formats that NumPy still does not model well. The test suite compares `bfloat16`, `float8_e4m3`, `float8_e5m2`, `float6_e2m3`, and `float6_e3m2` behavior against `ml_dtypes` where that comparison is meaningful. Promotion is intentional. Mixed exotic floats are routed through wider compute types rather than pretending a same-width accumulator is good enough. ### Buffer-First Casting `nk.astype` converts a CPU buffer in one native pass. Unlike `nk.Tensor(a).astype(dtype)`, it never stages `a` through an intermediate NumKong `Tensor`. The input may have arbitrary rank and strides, and the returned `Tensor` preserves its shape and is C-contiguous. ```python image_u8 = np.random.randint(0, 256, size=(1080, 1920, 3), dtype=np.uint8) image_f32 = nk.astype(image_u8, "float32") # Reuse an existing allocation — writes in place and hands `out` back. out = np.empty(image_u8.shape, dtype=np.float32) assert nk.astype(image_u8, "float32", out=out) is out ``` Any writable buffer works as `out`, whether a NumPy array, a `memoryview`, or a NumKong `Tensor`. It must be C-contiguous, exactly the input shape, already of the requested dtype, and must not overlap the input. A NumKong `Tensor` is the way to name a dtype NumPy cannot express, such as `bfloat16`, `uint1`, `int4`, or `uint4`. The native conversion runs without the GIL. Packed dtypes — `uint1`, `int4`, `uint4`, and `e2m1` — count __logical dimensions__ in a `Tensor` shape. Dimension count must be a multiple of the values per byte: 8 for `uint1` and 2 for the nibble types, with element 0 in the high nibble. The buffer protocol, `__array_interface__`, `nbytes`, and DLPack expose the whole bytes, so a raw NumPy byte array passed with `dtype="int4"` holds two dimensions per byte. ```python codes = np.arange(16, dtype=np.uint8) % 8 packed = nk.astype(codes, "int4") assert packed.shape == (16,) and packed.nbytes == 8 assert memoryview(packed).shape == (8,) # whole bytes assert nk.Tensor(np.asarray(memoryview(packed)), dtype="int4").shape == (16,) assert packed[3] == 3 and nk.dot(packed, packed) == int(codes.astype(np.int64) @ codes) ``` The conversion policy is NumKong's core policy. Same-dtype conversion makes an independent copy, and narrower floating formats use round-to-nearest, ties-to-even. Float-to-integer conversion rounds ties to even, saturates finite overflows and infinities, and maps `NaN` to zero. All NumKong dtype names are accepted as targets. ### ml_dtypes Interoperability NumKong accepts `ml_dtypes` arrays directly — no `.view(np.uint8)` workaround needed: ```python import ml_dtypes a = np.random.randn(100, 768).astype(np.float32).astype(ml_dtypes.bfloat16) b = np.random.randn(100, 768).astype(np.float32).astype(ml_dtypes.bfloat16) result = nk.cdist(a, b, "dot") # just works ``` NumKong scalars also work as NumPy dtype specifiers: ```python arr = np.array([1.0, 2.0, 3.0], dtype=nk.bfloat16) float(arr[0]) # → 1.0 ``` Type name mapping between the two libraries: | ml_dtypes | NumKong | Status | | :----------------------------- | :--------------------------- | :---------------------------------------- | | `ml_dtypes.bfloat16` | `nk.bfloat16` / `"bfloat16"` | Identical format | | `ml_dtypes.float8_e4m3` | `nk.float8_e4m3` / `"e4m3"` | Identical (IEEE E4M3) | | `ml_dtypes.float8_e4m3fn` | `nk.float8_e4m3` / `"e4m3"` | Identical (E4M3FN = no inf) | | `ml_dtypes.float8_e5m2` | `nk.float8_e5m2` / `"e5m2"` | Identical format | | `ml_dtypes.float6_e2m3fn` | `nk.float6_e2m3` / `"e2m3"` | Identical (MX E2M3) | | `ml_dtypes.float6_e3m2fn` | `nk.float6_e3m2` / `"e3m2"` | Identical (MX E3M2) | | `ml_dtypes.float8_e4m3fnuz` | … | Rejected: different bias, NaN, and zero | | `ml_dtypes.float8_e5m2fnuz` | … | Rejected: different NaN and zero encoding | | `ml_dtypes.float8_e4m3b11fnuz` | … | Rejected: bias=11, incompatible encoding | | `ml_dtypes.float8_e8m0fnu` | `"float8_e8m0"` / `"e8m0"` | MX block scale byte (see `ScaledTensor`) | | `ml_dtypes.float8_e3m4` | … | Not supported: no NumKong kernel | | `ml_dtypes.float4_e2m1fn` | `"float4_e2m1"` / `"e2m1"` | NVFP4/MXFP4 element (see `ScaledTensor`) | | `ml_dtypes.int4` | `"int4"` | Compatible via buffer protocol | | `ml_dtypes.uint4` | `"uint4"` | Compatible via buffer protocol | | `ml_dtypes.int2` | … | Not supported | | `ml_dtypes.uint2` | … | Not supported | ### Block-Scaled Formats (OCP MX & NVIDIA NVFP4) Block-scaled formats group elements into fixed-size blocks, each carrying its own scale byte. | Format | Block | Element | Scale | Global | | :------------ | ----: | :------ | :------ | :-------- | | `"nvfp4"` | 16 | `e2m1` | `ue4m3` | `float32` | | `"mxfp4"` | 32 | `e2m1` | `ue8m0` | … | | `"mxfp6e2m3"` | 32 | `e2m3` | `ue8m0` | … | | `"mxfp6e3m2"` | 32 | `e3m2` | `ue8m0` | … | | `"mxfp8e4m3"` | 32 | `e4m3` | `ue8m0` | … | | `"mxfp8e5m2"` | 32 | `e5m2` | `ue8m0` | … | | `"mxint8"` | 32 | `i8` | `ue8m0` | … | Quantization happens through the same `astype` verb used for dense casts. Passing a block-scaled dtype name to a dense `Tensor.astype(...)` returns a `ScaledTensor`; calling `astype(...)` on a `ScaledTensor` with a dense dtype materializes it back. Quantization is always along the __last axis__, which must be a whole multiple of the block size. A `ScaledTensor` exposes plain read attributes, but no getter/setter methods: | Attribute | Type | Meaning | | :-------------- | :--------------- | :-------------------------------------------------------------------- | | `.elements` | `Tensor` | Element codes (e.g. `e2m1`), in the logical dense shape | | `.block_scales` | `Tensor` | One scale byte per block (`ue4m3` / `ue8m0`), last axis = `n / block` | | `.tensor_scale` | `float` / `None` | NVFP4 per-tensor multiplier; `None` for the MX family | | `.block_size` | `int` | Elements per block (16 for NVFP4, 32 for MX) | | `.dtype` | `str` | Composite format name (e.g. `"nvfp4"`) | | `.shape` | `tuple` | Logical (dense) shape | ```python import numkong, numpy as np x = np.random.randn(64).astype(np.float32) t = numkong.Tensor(x) # Encode → ScaledTensor. For NVFP4 the per-tensor scale is auto-derived and surfaced as `.tensor_scale`. q = t.astype("nvfp4") assert q.dtype == "nvfp4" and q.block_size == 16 assert isinstance(q.tensor_scale, float) # NVFP4 carries a per-tensor float assert q.block_scales.shape == (4,) # 64 / 16 blocks # Decode / materialize back to a dense Tensor of any standard dtype. out = q.astype("float32") # also "float16", "bf16", ... # MX formats carry no per-tensor scale. qm = t.astype("mxfp4") assert qm.tensor_scale is None and qm.block_size == 32 # Block-aligned slicing returns a ScaledTensor view sharing storage; both children slice in lockstep. X = numkong.Tensor(np.random.randn(3, 64).astype(np.float32)).astype("nvfp4") row = X[1] # → ScaledTensor of shape (64,) columns = X[:, 16:48] # → ScaledTensor of shape (3, 32); start/stop must be block-aligned materialized = row.astype("float32") ``` The `.elements` and `.block_scales` tensors are DLPack-exportable for zero-copy hand-off to PyTorch, CuPy, and friends. The `e2m1`, `e8m0`, and `ue4m3` bytes exchange as `kDLFloat4_e2m1fn`, `kDLFloat8_e8m0fnu`, and `kDLFloat8_e4m3fn`, with the capsule's last extent counting bytes: ```python import torch scales = torch.from_dlpack(q.block_scales) # zero-copy view of the scale bytes ``` ## Tensors and Views `Tensor` is a memoryview-backed object with NumPy-like metadata. It is the central container for strided views, transpose, reshape, flatten, and axis reductions. ```python import numpy as np import numkong as nk t = nk.Tensor(np.arange(12, dtype=np.float32).reshape(3, 4)) print(t.shape, t.dtype, t.ndim, t.strides, t.itemsize, t.nbytes) print(np.asarray(t)) # zero-copy array view when layout allows it print(t.T.shape) # transposed Tensor view print(t.reshape(2, 6).shape) print(t.flatten().shape) # Slicing — row, column, and scalar access row0 = t[0, :] # first row, shape (4,) col2 = t[:, 2] # third column, strided view, shape (3,) val = t[1, 2] # scalar element access → 6.0 # Reductions compose with sliced views idx = col2.argmin() # index of the minimum in the third column mn, i0, mx, i1 = col2.minmax() ``` The important layout rules are: - `Tensor` preserves shape and byte strides. - Transpose and slicing can produce non-contiguous views. - General reductions accept those views. - Matrix-style packed kernels require row-contiguous left operands. - Packed and symmetric outputs require C-contiguous `out` buffers. ### Memory Layout Requirements | API family | Input requirement | Output requirement | | :----------------------------------------------- | :------------------------------------------------------------------------------------------------- | :------------------------------------------------------------------------ | | Dense distances (`dot`, `euclidean`, etc.) | Rows must be contiguous (`strides[last] <= itemsize`). Strided rows (sliced columns) are rejected. | `out=` can have any stride along dim 0, but inner dim must be contiguous. | | `cdist` | Same as dense distances | `out=` must be rank-2 with shape `(a.count, b.count)` | | Elementwise (`scale`, `blend`, `fma`) | Arbitrary strides (strided views are supported) | `out=` must match input shape; strides are preserved | | Packed matrix (`dots_packed`) | Left operand: rank-2, contiguous rows, no negative strides | Output: C-contiguous with expected dtype | | Symmetric (`dots_symmetric`) | Contiguous rows | `out=`: C-contiguous square matrix | | Tensor reductions (`sum`, `min`, `argmin`, etc.) | Arbitrary strides (strided views supported) | N/A (returns scalar or reduced tensor) | ### GPU Tensors through DLPack Builds configured with `-C cmake.define.NUMKONG_CUDA_ARCH=all` or `-C cmake.define.NUMKONG_ROCM_ARCH=all` also carry the CUDA or ROCm kernels. `nk.cuda_count_devices()` counts the devices this process sees, zero without the runtime or its kernels, and `nk.cuda_capabilities_enabled(0)` reports the mask device 0 dispatches with, raising `ValueError` past the last. `nk.cuda_capabilities_detected` and `nk.cuda_capabilities_compiled` answer as their CPU twins do, `nk.rocm_*` and `nk.metal_*` do the same for the other vendors, and `cpu_configure_thread` is CPU only. `nk.from_dlpack` imports a CUDA or ROCm tensor without a copy and records its device. Only `dots_pack`, the dots, angulars and euclideans `_packed` and `_symmetric` functions, and `attention_pack`, `attention_packed` and `attention_packed_gradients` run on it, with that device's enabled capabilities. The shape functions also take it for their outputs, running the mask their packs were copied with. Element access, NumPy conversion, `cdist`, `maxsim` and every other CPU kernel raise `BufferError` rather than read device memory. No verb allocates, so every GPU call writes into an `out=` its device reaches. The dots and attention families take that as a Tensor on the same device, and mixing devices raises `ValueError`. The shape functions take any memory the device reaches, a Tensor or a buffer over managed or pinned memory, which the runtime checks per call, raising `BufferError` for the rest. Unified memory comes only from an explicit `copy`, as in `nk.shapes_pack(clouds).copy(capabilities=gpu)` and `nk.zeros(shape).copy(capabilities=gpu)`, which frees with the tensor and its views and refuses `resize()` and `reserve()`. A GPU call queues one launch on `stream=`, an integer handle defaulting to 0, the legacy stream, and returns before it finishes, as PyTorch does, and `nk.synchronize(capabilities, stream=stream)` waits for it. `nk.stream_init_cuda(ordinal)` makes such a handle on any device by its ordinal, and `nk.stream_free_cuda(stream)` frees it once synchronized, with `rocm_` and `metal_` twins. Metal tensors are host-readable and keep running on the CPU kernels, except the shape functions, which run on Metal over blocks of its allocator. ```python import torch import numkong as nk left = torch.randn(4096, 768, device="cuda", dtype=torch.bfloat16) right = torch.randn(8192, 768, device="cuda", dtype=torch.bfloat16) capabilities = nk.cuda_capabilities_enabled(right.device.index) packed_storage = torch.empty(nk.PackedMatrix.pack_size(8192, 768, "bf16", capabilities=capabilities), device="cuda", dtype=torch.uint8) result = torch.empty(4096, 8192, device="cuda", dtype=torch.float32) stream = torch.cuda.current_stream().cuda_stream packed = nk.dots_pack(nk.from_dlpack(right), out=nk.from_dlpack(packed_storage), stream=stream) nk.dots_packed(nk.from_dlpack(left), packed, out=nk.from_dlpack(result), stream=stream) nk.synchronize(capabilities, stream=stream) ``` ## Dot Products Dot products are their own family because storage type, conjugation rules, and output widening matter. ```python import numpy as np import numkong as nk a = (np.random.randn(256) + 1j * np.random.randn(256)).astype(np.complex64) b = (np.random.randn(256) + 1j * np.random.randn(256)).astype(np.complex64) dot = nk.dot(a, b) # numpy.dot(a, b) vdot = nk.vdot(a, b) # numpy.vdot(a, b) print(dot, vdot) ``` Real low-precision inputs can also be routed through explicit dtype tags when the storage buffer itself is raw bytes. ## Dense Distances 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. ```python 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. ### Output Control: `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. ```python 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 result ``` When `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: ```python nk.dot(a, b, dtype=nk.bfloat16) # works faster nk.dot(a, b, dtype="bfloat16") # works a bit slower ``` ## Set Similarity Packed-binary metrics operate on packed bits. That is why the right NumPy equivalent uses `np.packbits`, not `bool` arrays fed to scalar Python code. A `uint1` Tensor counts bits in its shape, while a raw byte array passed with `dtype="uint1"` holds 8 bits per byte. ```python import numpy as np import numkong as nk a_bits = np.random.randint(0, 2, size=256, dtype=np.uint8) b_bits = np.random.randint(0, 2, size=256, dtype=np.uint8) a, b = np.packbits(a_bits), np.packbits(b_bits) # 32 bytes each hamming = nk.hamming(a, b, dtype="uint1") jaccard = nk.jaccard(a, b, dtype="uint1") a_tensor = nk.astype(a_bits, "uint1") # shape (256,), 32 bytes of storage ``` Integer set Jaccard works on sorted ascending arrays of integer identifiers. Both inputs must be sorted in ascending order for correct results. ```python set_a = np.array([1, 3, 5, 7, 9], dtype=np.uint32) # must be sorted ascending set_b = np.array([3, 5, 8, 9, 10], dtype=np.uint32) # must be sorted ascending jaccard_sets = nk.jaccard(set_a, set_b) # |A ∩ B| / |A ∪ B| assert 0.0 < jaccard_sets < 1.0, "|A ∩ B| / |A ∪ B| should be in (0, 1)" ``` ## Probability Metrics Probability divergences deserve their own section because they are not just "one more distance". ```python import numpy as np import numkong as nk p = np.array([0.2, 0.3, 0.5], dtype=np.float32) q = np.array([0.1, 0.3, 0.6], dtype=np.float32) kl_forward, kl_reverse = nk.kullbackleibler(p, q), nk.kullbackleibler(q, p) assert kl_forward != kl_reverse, "KLD is asymmetric" js_forward, js_reverse = nk.jensenshannon(p, q), nk.jensenshannon(q, p) np.testing.assert_allclose(js_forward, js_reverse, atol=1e-6) # JSD is symmetric ``` ## Geospatial Metrics Geospatial kernels score every row point against every column point into a `[rows, columns]` matrix. Inputs are in radians, and outputs are in metres. Haversine and Vincenty take both sides as packs, built once per point set, so the many-to-many kernels evaluate no trigonometry before their per-pair work. ```python import numpy as np import numkong as nk # Statue of Liberty, Big Ben, Eiffel Tower latitudes = np.radians([40.6892, 51.5007, 48.8584]) longitudes = np.radians([-74.0445, -0.1246, 2.2945]) ellipsoid = nk.vincentys_pack(latitudes, longitudes) # any input dtype is cast to float64 sphere = nk.haversines_pack(latitudes, longitudes) assert len(sphere) == sphere.points == 3 and sphere.shape == (3, 2) vincentys = np.asarray(nk.vincentys_packed(ellipsoid, ellipsoid)) # Liberty → Big Ben ≈ 5,589,857 m on WGS-84 haversines = np.asarray(nk.haversines_packed(sphere, sphere)) # Liberty → Big Ben ≈ 5,574,848 m on the mean sphere assert vincentys.shape == haversines.shape == (3, 3) # Earth-fixed centimetres square exactly in 64 bits stations = np.array([[631_000_000, 0, 0], [0, 637_800_000, 0]], dtype=np.int32) squares = np.asarray(nk.sqeuclideans3d(stations, stations)) assert squares.dtype == np.uint64 and squares[0, 1] == 631_000_000**2 + 637_800_000**2 ``` `nk.haversines_packed` measures 2R · atan2(‖a − b‖, ‖a + b‖) with the IUGG mean radius of 6,371,008.8 m. That form stays conditioned from a millimetre to the antipode, and reads within 0.5% of the ellipsoidal geodesic. `nk.vincentys_packed` iterates Vincenty's inverse formula on the WGS-84 ellipsoid until the longitude settles within 1e-12 radians. A pair near the antipode that never settles reports its last iterate after 100 iterations. `nk.sqeuclideans3d` takes two `[n, 3]` matrices, both `int32` or both `float64`, and returns `uint64` or `float64` squares. The `int32` squares are exact whenever every axis difference fits 32 bits, which Earth-fixed centimetres below 4,359 km of altitude always satisfy. Its `threads=` argument splits the row points across workers. `nk.HaversinesPackedMatrix.pack_size(points)` and `nk.VincentysPackedMatrix.pack_size(points)` return the bytes a pack takes. A pack keeps its point count in a header ahead of its planar arrays, so a view of fewer points would need a header of its own. Indexing therefore returns the pack itself for a key selecting every point in order, and raises `ValueError` for any other. ## Curved Metrics Curved-space kernels use an extra metric tensor or inverse covariance and should not be mixed into the Euclidean section. ```python import numpy as np import numkong as nk # Complex bilinear form: aᴴ M b a = (np.ones(16) + 1j * np.zeros(16)).astype(np.complex64) b = (np.zeros(16) + 1j * np.ones(16)).astype(np.complex64) m = np.eye(16, dtype=np.complex64) bilinear = nk.bilinear(a, b, m) # Real Mahalanobis distance: √((a−b)ᵀ M⁻¹ (a−b)) x = np.ones(32, dtype=np.float32) y = np.full(32, 2.0, dtype=np.float32) inv_cov = np.eye(32, dtype=np.float32) mahalanobis = nk.mahalanobis(x, y, inv_cov) ``` ## Elementwise and Reductions Elementwise arithmetic and fused operations are their own family. They share the tensor infrastructure but should not be collapsed into the reduction or matrix sections. ```python import numpy as np import numkong as nk a = np.arange(8, dtype=np.float32) b = np.arange(8, dtype=np.float32)[::-1].copy() scaled = nk.scale(a, alpha=2.0, beta=1.0) # 2 * a + 1 blended = nk.blend(a, b, alpha=0.25, beta=0.75) fused = nk.fma(a, b, a, alpha=1.0, beta=1.0) # a * b + a assert np.asarray(scaled).shape == (8,) assert np.asarray(fused).shape == (8,) ``` ### Moments Reductions Moments reductions return `(sum, sum_of_squares)`. The key property is that NumKong does not force you into same-storage accumulation. ```python import numpy as np import numkong as nk x = np.full(4096, 255, dtype=np.uint8) nk_sum, nk_sumsq = nk.moments(nk.Tensor(x)) naive_sum = np.sum(x, dtype=np.uint8) # overflows immediately naive_sumsq = np.sum(x * x, dtype=np.uint8) # also overflows print(nk_sum, nk_sumsq, naive_sum, naive_sumsq) assert nk_sum > int(naive_sum) assert nk_sumsq > int(naive_sumsq) ``` Same-width accumulation is a bad default for low-precision storage. ### Min/Max Reductions Min/max reductions are in a separate section because they cover strided reduction cases. NumKong provides SIMD-accelerated strided reductions that are not common in other libraries. ```python import numpy as np import numkong as nk matrix = nk.Tensor(np.array([ [ 3.0, 0.0, 7.0], [ 1.0, 2.0, 5.0], [ 4.0, -1.0, 6.0], ], dtype=np.float32)) second_column = matrix[:, 1] # strided view into a row-major Nx3 tensor idx = second_column.argmin() mn, i0, mx, i1 = second_column.minmax() assert idx == 2 assert int(i0) == 2 assert float(np.asarray(mn)) == -1.0 ``` Fresh measurement for the rewritten docs: on an Apple M2 Pro, `np.argmin(matrix[:, 1])` on a row-major `2,000,000 x 3` `float32` array took about `1.63 ms` median. The equivalent NumKong `Tensor(... )[:, 1].argmin()` took about `0.67 ms` median. That is about `2.45x` faster on this strided reduction case. ## Sparse Operations Sparse helpers cover both sorted-index intersections and weighted sparse dot products. ```python import numpy as np import numkong as nk idx_a, idx_b = np.array([1, 3, 5, 7], dtype=np.uint32), np.array([3, 4, 5, 8], dtype=np.uint32) intersection_size = nk.intersect(idx_a, idx_b) # len(np.intersect1d(idx_a, idx_b)) assert intersection_size == 2, "indices 3 and 5" val_a, val_b = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), np.array([5.0, 6.0, 7.0, 8.0], dtype=np.float32) sparse_dot = nk.sparse_dot(idx_a, val_a, idx_b, val_b) assert sparse_dot > 0, "weighted dot over shared indices" ``` ## Matrix Kernels Packed matrix kernels are the right tool when the right-hand side is reused across many query batches. This is the `GEMM`-like story. ```python import numpy as np import numkong as nk left = np.random.randn(128, 768).astype(np.float32) right = np.random.randn(10_000, 768).astype(np.float32) right_packed = nk.dots_pack(right, dtype="float32") # pack once, reuse many times scores = nk.dots_packed(left, right_packed) # equivalent to left @ right.T assert scores.shape == (128, 10_000) assert right_packed.nbytes == nk.PackedMatrix.pack_size(10_000, 768, dtype="float32") ``` Important runtime rules from the current implementation: - `a` must be rank-2 - `a` must have contiguous rows - negative strides are rejected for these matrix kernels - `out`, when provided, must be C-contiguous with the expected dtype - `rows_begin` and `rows_end` split the left operand rows The arithmetic advantages are: - one-time packing of `B` - one-time internal layout conversion and depth padding - norm reuse for `angulars_packed` and `euclideans_packed` - no repeated scan of the original right-hand-side layout Packing itself does not require aligned caller buffers. The packed object owns its internal payload and handles the layout under the hood. `Tensor @ PackedMatrix` is also supported and maps to the same packed dot-product path. ### All-Pairs APIs and cdist `cdist` is the NumPy/SciPy-shaped all-pairs entrypoint. It handles rectangular matrix pairs and symmetric self-distance cases. ```python import numpy as np import numkong as nk queries = np.random.randn(100, 768).astype(np.float32) database = np.random.randn(10_000, 768).astype(np.float32) pairwise = nk.angular(queries, database[:100]) # rectangular broadcasted pairwise call all_pairs = nk.cdist(queries, database, metric="angular") # scipy.spatial.distance.cdist analogue assert np.asarray(pairwise).shape == (100, 100) assert np.asarray(all_pairs).shape == (100, 10_000) ``` `cdist` and the packed and symmetric kernels take a `threads=` argument, where `0` means every logical processor, and row ranges are still available for partitioning the work yourself. `dots_pack` and `hammings_pack` take it too, splitting the packed matrix by columns. ### Symmetric Kernels for SYRK-Like Workloads Symmetric kernels solve a different problem from packed cross-matrix kernels. They compute self-similarity or self-distance matrices. This is the `SYRK`-like story. ```python import numpy as np import numkong as nk vectors = np.random.randn(1024, 768).astype(np.float32) out = nk.zeros((1024, 1024), dtype="float64") nk.dots_symmetric(vectors, out=out, rows_begin=0, rows_end=256) nk.dots_symmetric(vectors, out=out, rows_begin=256, rows_end=512) assert out.shape == (1024, 1024) ``` This family has different economics from packed `GEMM`-like work. It avoids duplicate `(i, j)` and `(j, i)` evaluations. It is naturally partitioned by row windows of one square output. `angulars_symmetric` and `euclideans_symmetric` also benefit from reuse of dot-product-derived work inside the symmetric sweep. That is why these APIs are faster than a nested Python loop over `angular(a[i], a[j])`. ### Attention with a Pre-Packed KV-Cache Ragged scaled-dot-product attention packs the key/value tokens once, then scores many queries against them. Segments of different lengths share one packed cache, addressed by cumulative `key_offsets`. A segment may hold fewer keys than its slot, as a decode cache with spare room does: pass `key_lengths` too, and the keys past each length are left out of the pack. ```python import numpy as np import numkong as nk lengths = [6, 10] # two ragged segments offsets = np.array([0, *np.cumsum(lengths)], dtype=np.uint32) tokens, heads, depth = int(offsets[-1]), 2, 64 k = nk.Tensor(np.random.randn(tokens, heads * depth).astype(np.float32)).astype("bfloat16") v = nk.Tensor(np.random.randn(tokens, heads * depth).astype(np.float32)).astype("bfloat16") kv = nk.attention_pack(k, v, key_offsets=offsets, depth=depth) queries = nk.Tensor(np.random.randn(tokens, heads * depth).astype(np.float32)).astype("bfloat16") out = nk.attention_packed(queries, kv, query_offsets=offsets) windowed = nk.attention_packed(queries, kv, query_offsets=offsets, keys_before=3, keys_after=0) assert kv.shape == (kv.key_value_head_count, kv.depth, kv.segments) ``` Queries align to the end of each segment's keys: row `r` of `q` queries over `k` keys sits at position `r + k - q` and sees `keys_before` keys before it and `keys_after` after, `None` being unbounded. So `keys_after=0` is causal for prefill and decode alike, and adding `keys_before=w - 1` slides a window of `w` keys. Rows that see no key come back as zeros. For training, the forward also fills a `(tokens, heads)` log-sum-exp, and `attention_packed_gradients` turns it, the output and its gradient into BF16 query, key and value gradients under the same band. CPUs compute them serially, so the cache is packed with `nk.Capability.SERIAL`: ```python kv = nk.attention_pack(k, v, key_offsets=offsets, depth=depth, capabilities=nk.Capability.SERIAL) log_sum_exp = np.empty((tokens, heads), dtype=np.float32) out = nk.attention_packed(queries, kv, query_offsets=offsets, keys_after=0, log_sum_exp=log_sum_exp) output_gradient = np.random.randn(tokens, heads * depth).astype(np.float32) query_gradient, key_gradient, value_gradient = nk.attention_packed_gradients( queries, kv, query_offsets=offsets, output=out, output_gradient=output_gradient, log_sum_exp=log_sum_exp, keys_after=0) ``` `query_gradient=` takes a pre-allocated `(tokens, heads * depth)` f32 tensor, whose rows may sit in a wider buffer. With `threads=`, attention cuts its grid of query tokens × heads into a few windows of equal cost per thread, so a single decode token still spreads over its heads. ## MaxSim 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") ``` ## Mesh Alignment ### Many-to-Many Shape Distances Point clouds are packed once each, centered on their centroids and rotated onto their principal axes. The packed sets then score every row cloud against every column cloud, and the clouds may differ in size. Chamfer and Hausdorff compare nearest partners under the 8 sign flips of the principal axes. Wasserstein-Procrustes solves for the exact transport plan and the best rotation or reflection together. On equal sizes it is at most the Kabsch RMSD of a known correspondence, and equals it when that correspondence is an optimal plan. Gromov-Wasserstein compares the clouds' own distance matrices. ```python import numpy as np import numkong as nk molecules = [np.random.randn(atoms, 3) for atoms in (20, 31, 60)] rows = nk.shapes_pack(molecules) # float64 input packs as float32 columns = nk.shapes_pack(molecules[1:]) assert rows.shape == (3, 111) and columns.shape == (2, 91) # clouds and total points rotations = np.empty((3, 2, 3, 3), dtype=np.float32) distances = nk.wassersteinprocrustes_packed(rows, columns, rotations=rotations) assert np.asarray(distances).shape == (3, 2) assert np.asarray(distances)[1, 0] < 1e-6, "a molecule against itself" chamfers = nk.chamfers_packed(rows[:1], columns) # slices are views over the same packs assert np.asarray(chamfers).dtype == np.float64 assert np.allclose(np.asarray(rows.unpack(1)), molecules[1], atol=1e-5) # the inverse of packing grown = nk.shapes_pack_normalized([molecules[0] * 3.7 + 1.0]) # divided by its RMS radius originals = nk.shapes_pack_normalized(molecules) assert np.asarray(nk.wassersteinprocrustes_packed(originals, grown))[0, 0] < 1e-5 ``` Each rotation takes column cloud `b` onto row cloud `a` through the centroids: `a ≈ (b - b.mean(0)) @ R.T + a.mean(0)`. The optional `rotations=` output is a writable C-contiguous `float32` buffer of shape `[rows, columns, 3, 3]`, and the functions always return the distances. `nk.shapes_pack_normalized` also divides every cloud by its RMS distance from the centroid, so every kind ignores size as well as pose. Between normalized packs, the rotation also scales by `scale_a / scale_b`, and each kind compares the packed coordinates, so mixing absolute and normalized packs in one call is the caller's choice. The pack layout is private, and the `centroids`, `frames`, `scales`, `spreads` and `max_radii` properties report every cloud's pose and bounds. `unpack(index)` returns one cloud as a `[points, 3]` tensor in the pack dtype and its original pose and scale, within one storage rounding. `nk.ShapesPackedMatrix.pack_size(points, dtype)` returns the bytes one pack takes. Chamfer and Hausdorff take `limit=`: distances up to it are exact, and a pair past it may stop early and report a lower bound that also lies past it. Chamfer, Wasserstein-Procrustes and Gromov-Wasserstein allocate their workspace inside the call, or take a writable `workspace=` of `nk.chamfers_workspace_size(a, b, threads=threads)` bytes, with its `wassersteinprocrustes` and `gromovwassersteins` siblings. `nk.rmsds_packed(a, b)` scores clouds of equally many points, point i matching point i under the best rotation. Every kind takes `out=`, a writable `[rows, columns]` buffer of its output dtype. Every kind takes `threads=`, splitting the row clouds with a workspace per worker, and both packing functions take it to split the clouds. Packs need no GPU kernel: the host gathers them, and `copy` places the table and every pack in one block of unified memory on the GPU the mask names. Every shape function then runs that mask in one launch on `stream=`, ignoring `threads=`, so it needs `out=` and, where the workspace size is not zero, a single window of `workspace=`. ```python gpu = nk.cuda_capabilities_enabled(0) packs = nk.shapes_pack(clouds, dtype="f16", threads=8).copy(capabilities=gpu) distances = nk.from_dlpack(torch.empty((n, n), device="cuda")) nk.chamfers_packed(packs, packs, out=distances, limit=2.5, stream=stream) nk.synchronize(gpu, stream=stream) ``` A slice of copied packs in steps would rebuild the table on the host, so it raises `BufferError` on the GPU until it is copied itself. Metal runs `float16` packs only, and writes only into blocks of its own allocator, so a plain NumPy `out=` raises `BufferError`: ```python gpu = nk.metal_capabilities_enabled(0) packs = nk.shapes_pack(clouds, dtype="f16").copy(gpu) distances = nk.zeros((n, n), dtype="float32").copy(capabilities=gpu) nk.chamfers_packed(packs, packs, out=distances, limit=2.5) nk.synchronize(gpu) ``` ## Parallelism Capability detection is explicit, through module functions named after each device group: ```python import numpy as np import numkong as nk # `enabled` is what dispatch runs within: detected on this CPU AND compiled into the wheel, settled as the library loads. print(repr(nk.cpu_capabilities_enabled())) print(nk.Capability.SKYLAKE in nk.cpu_capabilities_enabled()) # will AVX-512 kernels run here? # The two raw axes, when you specifically mean one of them: print(repr(nk.cpu_capabilities_detected())) # this CPU print(repr(nk.cpu_capabilities_compiled())) # this build # Prepare a thread that runs kernels, like Arm's FPCR; the library already asked for AMX tiles as it loaded. nk.cpu_configure_thread(nk.cpu_capabilities_enabled()) # Narrow a single call, e.g. to test one capability: what cannot run here is dropped. a = np.random.randn(1536).astype(np.float32) print(nk.dot(a, a, capabilities=nk.Capability.SERIAL)) # the reference kernel ``` Every function that runs kernels takes the mask as `capabilities=` and a GPU stream as `stream=`, an integer pointer left `None` on the CPU. It defaults to every CPU capability, which the library clamps to the enabled ones, so the binding keeps no mask of its own. That mask keeps only CPU capabilities, as GPU tensors dispatch with their own device's, see [GPU Tensors through DLPack](#gpu-tensors-through-dlpack). Packed operands, like `PackedMatrix`, remember the mask that packed them and are read with it, since each capability lays its packs out differently. The mask model is the one in [Dispatch Points & Capability Masks](https://github.com/ashvardanian/NumKong/blob/main/README.md#dispatch-points--capability-masks). The current implementation releases the GIL around the native dense metric calls and around the packed and symmetric matrix kernels. The repository also has threading tests for packed and symmetric row-range partitioning. `GEMM`-like packed work and `SYRK`-like symmetric work should be documented differently: ```python import concurrent.futures import numpy as np import numkong as nk left = np.random.randn(4096, 768).astype(np.float32) right = np.random.randn(8192, 768).astype(np.float32) packed = nk.dots_pack(right, dtype="float32") out = nk.zeros((4096, 8192), dtype="float64") # out must be pre-allocated with correct shape and dtype def packed_chunk(start, end): nk.dots_packed(left, packed, out=out, rows_begin=start, rows_end=end) # split left rows against one shared packed RHS with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: for start in range(0, 4096, 1024): pool.submit(packed_chunk, start, min(start + 1024, 4096)) ``` ```python import concurrent.futures import numpy as np import numkong as nk vectors = np.random.randn(4096, 768).astype(np.float32) out = nk.zeros((4096, 4096), dtype="float64") # out must be pre-allocated with correct shape and dtype def symmetric_chunk(start, end): nk.dots_symmetric(vectors, out=out, rows_begin=start, rows_end=end) # split row windows of one square output with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: for start in range(0, 4096, 1024): pool.submit(symmetric_chunk, start, min(start + 1024, 4096)) ``` OpenMP and other native schedulers still matter in lower layers. For Python, the intended user-facing story is external partitioning around the GIL-free kernels you actually use. ## External Memory NumKong implements the Python buffer protocol for zero-copy interop with NumPy, PyTorch, and other buffer-aware libraries. Two additional primitives cover pointer-level workflows: `data_ptr` reads the integer address out of any `Tensor`, and `from_pointer()` wraps any integer address back into one. `data_ptr` returns the raw address, suitable for passing into ctypes, CUDA, or any FFI boundary. `from_pointer(address, shape, dtype, *, strides=None, owner=None)` creates a non-owning `Tensor` view. The optional `owner` keeps the source object alive for the lifetime of the view. ```python import numpy as np import numkong as nk # Round-trip through an integer address matrix = nk.zeros((3, 4), dtype='float32') address = matrix.data_ptr matrix_view = nk.from_pointer(address, (3, 4), 'float32', owner=matrix) # Wrap a NumPy array with zero copies embeddings = np.random.randn(1024).astype(np.float32) embeddings_view = nk.from_pointer(embeddings.ctypes.data, (1024,), 'float32', owner=embeddings) nk.dot(embeddings, embeddings_view) # same underlying data ``` PyTorch tensors already implement the buffer protocol, so most functions accept them directly. For explicit pointer-level control, or to go the other direction, the same primitives apply: ```python import torch query = torch.randn(512) nk.dot(query, query) # buffer protocol, zero copy # Explicit pointer wrap query_view = nk.from_pointer(query.data_ptr(), tuple(query.shape), 'float32', owner=query) # NumKong → PyTorch: 1D via buffer protocol, N-D via NumPy bridge flat = torch.frombuffer(memoryview(nk_tensor), dtype=torch.float32) shaped = torch.as_tensor(np.asarray(nk_tensor)) ``` CUDA unified memory, pinned buffers, and mmap'd files all work the same way — any CPU-accessible pointer is valid. ```python import ctypes, mmap # CUDA unified memory (ensure CPU accessibility first) cudart = ctypes.CDLL("libcudart.so") unified_ptr = ctypes.c_void_p() cudart.cudaMallocManaged(ctypes.byref(unified_ptr), 4096, 1) cudart.cudaDeviceSynchronize() unified = nk.from_pointer(unified_ptr.value, (1024,), 'float32') # Memory-mapped file with open("data.bin", "r+b") as f: mapping = mmap.mmap(f.fileno(), 0) mapped = nk.from_pointer(ctypes.addressof( ctypes.c_char.from_buffer(mapping)), (1024,), 'float32', owner=mapping) ```