usearch 2.25.1 → 2.26.0

package/numkong/test/test_cross_amx.cpp

@@ -46,5 +46,18 @@ void test_cross_amx() {
           nk_dots_pack_f16_graniteamx, nk_euclideans_packed_f16_graniteamx);
     check("euclideans_symmetric_f16_graniteamx", test_euclideans_symmetric<f16_t>,
           nk_euclideans_symmetric_f16_graniteamx);
+
+    check("dots_packed_e5m2_graniteamx", test_dots_packed<e5m2_t>, nk_dots_packed_size_e5m2_graniteamx,
+          nk_dots_pack_e5m2_graniteamx, nk_dots_packed_e5m2_graniteamx);
+    check("dots_symmetric_e5m2_graniteamx", test_dots_symmetric<e5m2_t>, nk_dots_symmetric_e5m2_graniteamx);
+
+    check("angulars_packed_e5m2_graniteamx", test_angulars_packed<e5m2_t>, nk_dots_packed_size_e5m2_graniteamx,
+          nk_dots_pack_e5m2_graniteamx, nk_angulars_packed_e5m2_graniteamx);
+    check("angulars_symmetric_e5m2_graniteamx", test_angulars_symmetric<e5m2_t>, nk_angulars_symmetric_e5m2_graniteamx);
+
+    check("euclideans_packed_e5m2_graniteamx", test_euclideans_packed<e5m2_t>, nk_dots_packed_size_e5m2_graniteamx,
+          nk_dots_pack_e5m2_graniteamx, nk_euclideans_packed_e5m2_graniteamx);
+    check("euclideans_symmetric_e5m2_graniteamx", test_euclideans_symmetric<e5m2_t>,
+          nk_euclideans_symmetric_e5m2_graniteamx);
 #endif
 }

package/numkong/test/test_dlpack.py

@@ -0,0 +1,492 @@
+#!/usr/bin/env python3
+"""DLPack zero-copy interop tests for NumKong.
+
+Covers the Python Array API DLPack protocol implemented in ``python/dlpack_interop.{c,h}`` and
+``python/dlpack_abi.h``. Tests are organized into per-framework sections; every test gates its dependencies
+through ``pytest.importorskip`` so the suite runs against any subset of installed frameworks (NumPy itself is
+treated as optional — only the protocol-contract tests at the bottom of this file are framework-free, and they
+build their tensors with NumKong's own constructors).
+
+Frameworks exercised:
+
+- NumPy           — CPU round-trips for f32 / f64 / f16 / i* / u*.
+- PyTorch         — CPU round-trips, plus bf16, float8_e4m3fn, float8_e5m2 dtype-fidelity tests.
+- JAX             — CPU round-trips for f32 / f64 / bf16.
+- TensorFlow      — CPU round-trips for f32 / f64 / bf16 via ``tf.experimental.dlpack``.
+- CuPy            — three guarded paths covering the importer's device-acceptance contract:
+                    (a) plain ``cudaMalloc`` (kDLCUDA) → must be REJECTED with a clear ValueError;
+                    (b) ``cudaMallocManaged`` unified memory (kDLCUDAManaged) → must be ACCEPTED;
+                    (c) ``cudaMallocHost`` pinned host memory (kDLCUDAHost) → must be ACCEPTED.
+- PyArrow         — Arrow Array → NumKong (one-way; PyArrow is producer-only in the protocol).
+- ONNX Runtime    — ``OrtValue`` ↔ NumKong via the ``to_dlpack`` / ``ortvalue_from_numpy`` helpers.
+- MLX             — Apple Silicon's unified-memory framework. Tests skip on Linux (no ``mlx.core`` wheel) and
+                    run on macOS contributor machines / Apple Silicon CI.
+
+The importer accepts every DLPack ``device_type`` whose pointer is dereferenceable from host code
+(kDLCPU / CUDAHost / ROCMHost / CUDAManaged / OneAPI / Metal). Pure device memory (kDLCUDA, kDLROCM,
+kDLOpenCL, kDLVulkan, kDLWebGPU, kDLHexagon, kDLMAIA, kDLTrn, kDLVPI, kDLExtDev) is rejected.
+"""
+
+import pytest
+
+import numkong as nk
+
+
+# region Helpers
+
+
+_INTEGER_DTYPES = ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]
+_PLAIN_FLOAT_DTYPES = ["float16", "float32", "float64"]
+_PLAIN_NUMERIC_DTYPES = _PLAIN_FLOAT_DTYPES + _INTEGER_DTYPES
+
+# Per-framework dtype matrices. Each entry is a `(framework_attr_name, nk_dtype_string)` pair so a single
+# parametrized test can drive both directions: `getattr(framework_module, attr)` resolves the producer-side
+# dtype, and the `nk_dtype_string` is what we expect back on the NumKong-side after import.
+#
+# We include the "weird GPU formats" (bf16, fp8 e4m3fn, fp8 e5m2, fp6 e2m3, fp6 e3m2) wherever the framework
+# supports them natively, so each framework's coverage matches what dtype-fidelity DLPack actually buys.
+
+# NumPy core has no bf16/fp8/fp6 — those would need ml_dtypes. Skip exotics here; they're tested via PyTorch/MLX.
+_NUMPY_DTYPES = [(name, name) for name in _PLAIN_NUMERIC_DTYPES]
+
+# PyTorch ships native bf16 + float8_e4m3fn + float8_e5m2 since 2.1+; full coverage here is the marquee case.
+_TORCH_DTYPES = (
+    [(name, name) for name in _PLAIN_NUMERIC_DTYPES]
+    + [("bfloat16", "bfloat16"), ("float8_e4m3fn", "e4m3"), ("float8_e5m2", "e5m2")]
+)
+
+# JAX has native bf16 (always) and the integer / float subset; no fp8/fp6 yet.
+_JAX_DTYPES = [
+    ("float16", "float16"), ("float32", "float32"), ("float64", "float64"),
+    ("int8", "int8"), ("int16", "int16"), ("int32", "int32"), ("int64", "int64"),
+    ("uint8", "uint8"), ("uint16", "uint16"), ("uint32", "uint32"), ("uint64", "uint64"),
+    ("bfloat16", "bfloat16"),
+]
+
+# TensorFlow: native bf16 + float16/32/64 + integers. No public fp8 dtype that maps cleanly via DLPack today.
+_TF_DTYPES = [
+    ("float16", "float16"), ("float32", "float32"), ("float64", "float64"),
+    ("int8", "int8"), ("int16", "int16"), ("int32", "int32"), ("int64", "int64"),
+    ("uint8", "uint8"), ("uint16", "uint16"), ("uint32", "uint32"), ("uint64", "uint64"),
+    ("bfloat16", "bfloat16"),
+]
+
+# PyArrow exposes float / signed-int / unsigned-int builders by name. No bf16/fp8/fp6.
+_PYARROW_DTYPES = [
+    ("float16", "float16"), ("float32", "float32"), ("float64", "float64"),
+    ("int8", "int8"), ("int16", "int16"), ("int32", "int32"), ("int64", "int64"),
+    ("uint8", "uint8"), ("uint16", "uint16"), ("uint32", "uint32"), ("uint64", "uint64"),
+]
+
+# MLX ships native bf16 on Apple Silicon, plus float16/32/integers. No fp8 yet (mlx.core 0.x).
+_MLX_DTYPES = [
+    ("float16", "float16"), ("float32", "float32"),
+    ("int8", "int8"), ("int16", "int16"), ("int32", "int32"), ("int64", "int64"),
+    ("uint8", "uint8"), ("uint16", "uint16"), ("uint32", "uint32"), ("uint64", "uint64"),
+    ("bfloat16", "bfloat16"),
+]
+
+
+def _assert_round_trip_through(producer_array, expected_dtype_str: str, expected_shape: tuple) -> None:
+    """Common pattern: framework-owned array → NumKong view → assertions on dtype/shape."""
+    imported = nk.from_dlpack(producer_array)
+    assert imported.shape == expected_shape, f"{imported.shape} != {expected_shape}"
+    assert imported.dtype == expected_dtype_str, f"{imported.dtype} != {expected_dtype_str}"
+
+
+# endregion Helpers
+
+
+# region NumPy round-trip
+
+
+@pytest.mark.parametrize("np_dtype, nk_dtype", _NUMPY_DTYPES)
+def test_numpy_export(np_dtype, nk_dtype):
+    """A NumKong tensor must be consumable zero-copy by `np.from_dlpack`."""
+    np = pytest.importorskip("numpy")
+    source = np.arange(24, dtype=np_dtype).reshape(4, 6)
+    tensor = nk.Tensor(source)
+    arr = np.from_dlpack(tensor)
+    assert arr.shape == (4, 6)
+    assert str(arr.dtype) == np_dtype
+    assert tensor.dtype == nk_dtype
+    np.testing.assert_array_equal(arr, source)
+    arr[0, 0] = 99
+    assert np.asarray(tensor)[0, 0] == 99, "mutation through NumPy view not visible in NumKong tensor"
+
+
+@pytest.mark.parametrize("np_dtype, nk_dtype", _NUMPY_DTYPES)
+def test_numpy_import(np_dtype, nk_dtype):
+    """A NumPy array must be consumable zero-copy by `nk.from_dlpack`."""
+    np = pytest.importorskip("numpy")
+    source = np.arange(20, dtype=np_dtype)
+    tensor = nk.from_dlpack(source)
+    assert tensor.shape == (20,)
+    assert tensor.dtype == nk_dtype
+    np.testing.assert_array_equal(np.asarray(tensor), source)
+
+
+def test_numpy_import_strided():
+    """A non-contiguous (transposed) NumPy source preserves strides on import."""
+    np = pytest.importorskip("numpy")
+    source = np.arange(24, dtype=np.float32).reshape(4, 6)
+    tensor = nk.from_dlpack(source.T)
+    assert tensor.shape == (6, 4)
+    np.testing.assert_array_equal(np.asarray(tensor), source.T)
+
+
+# endregion NumPy round-trip
+
+
+# region PyTorch round-trip
+
+
+@pytest.mark.parametrize("torch_dtype_attr, nk_dtype", _TORCH_DTYPES)
+def test_torch_import(torch_dtype_attr, nk_dtype):
+    """Every torch-supported dtype (including bf16, fp8 e4m3fn, fp8 e5m2) imports zero-copy as NumKong dtype."""
+    torch = pytest.importorskip("torch")
+    torch_dtype = getattr(torch, torch_dtype_attr)
+    # `torch.zeros` works for fp8 (where `torch.arange` is unsupported); the test only checks dtype + round-trip.
+    src = torch.zeros(4, 6, dtype=torch_dtype)
+    nk_tensor = nk.from_dlpack(src)
+    assert nk_tensor.shape == (4, 6)
+    assert nk_tensor.dtype == nk_dtype
+    pt_back = torch.from_dlpack(nk_tensor)
+    assert pt_back.dtype == torch_dtype
+    assert tuple(pt_back.shape) == (4, 6)
+
+
+@pytest.mark.parametrize("torch_dtype_attr, nk_dtype", _TORCH_DTYPES)
+def test_torch_export(torch_dtype_attr, nk_dtype):
+    """Every torch-supported dtype is consumable by `torch.from_dlpack` after a NumKong round-trip."""
+    torch = pytest.importorskip("torch")
+    torch_dtype = getattr(torch, torch_dtype_attr)
+    # Build the source on the torch side (avoids needing ml_dtypes for fp8/bf16 numpy paths).
+    src = torch.zeros(4, 6, dtype=torch_dtype)
+    nk_tensor = nk.from_dlpack(src)  # NumKong view sharing torch's buffer
+    pt_view = torch.from_dlpack(nk_tensor)  # back to torch via NumKong's exporter
+    assert pt_view.dtype == torch_dtype
+    assert tuple(pt_view.shape) == (4, 6)
+
+
+def test_torch_export_mutation_visibility():
+    """Writes through one view are visible through the other (proves zero-copy, not a value comparison)."""
+    np = pytest.importorskip("numpy")
+    torch = pytest.importorskip("torch")
+    source = np.arange(24, dtype=np.float32).reshape(4, 6)
+    tensor = nk.Tensor(source)
+    pt = torch.from_dlpack(tensor)
+    pt[0, 0] = 99
+    assert np.asarray(tensor)[0, 0] == 99, "mutation through PyTorch view not visible in NumKong tensor"
+
+
+def test_torch_strided_import():
+    """A transposed PyTorch tensor preserves strides when imported."""
+    np = pytest.importorskip("numpy")
+    torch = pytest.importorskip("torch")
+    src = torch.arange(24, dtype=torch.float32).reshape(4, 6).T
+    nk_tensor = nk.from_dlpack(src)
+    assert nk_tensor.shape == (6, 4)
+    np.testing.assert_array_equal(np.asarray(nk_tensor), src.numpy())
+
+
+# endregion PyTorch round-trip
+
+
+# region JAX round-trip
+
+
+@pytest.mark.parametrize("jax_dtype_attr, nk_dtype", _JAX_DTYPES)
+def test_jax_export(jax_dtype_attr, nk_dtype):
+    """NumKong → JAX zero-copy across every JAX-supported dtype incl. bf16."""
+    np = pytest.importorskip("numpy")
+    jax = pytest.importorskip("jax")
+    # JAX disables float64 by default and silently downcasts to float32 unless x64 is enabled.
+    if jax_dtype_attr in ("float64", "int64", "uint64"):
+        jax.config.update("jax_enable_x64", True)
+    jnp = jax.numpy
+    np_dtype = jax_dtype_attr if jax_dtype_attr != "bfloat16" else "float32"
+    source = np.arange(20, dtype=np_dtype)
+    tensor = nk.Tensor(source) if jax_dtype_attr != "bfloat16" else nk.Tensor(source).astype("bfloat16")
+    jax_array = jnp.from_dlpack(tensor)
+    assert jax_array.shape == (20,)
+    assert str(jax_array.dtype) == jax_dtype_attr
+    assert tensor.dtype == nk_dtype
+
+
+@pytest.mark.parametrize("jax_dtype_attr, nk_dtype", _JAX_DTYPES)
+def test_jax_import(jax_dtype_attr, nk_dtype):
+    """JAX → NumKong zero-copy across every JAX-supported dtype incl. bf16."""
+    jax = pytest.importorskip("jax")
+    if jax_dtype_attr in ("float64", "int64", "uint64"):
+        jax.config.update("jax_enable_x64", True)
+    jnp = jax.numpy
+    jax_dtype = getattr(jnp, jax_dtype_attr)
+    # `jnp.zeros` works for every supported dtype incl. bf16 / int / uint.
+    src = jnp.zeros(20, dtype=jax_dtype)
+    nk_tensor = nk.from_dlpack(src)
+    assert nk_tensor.shape == (20,)
+    assert nk_tensor.dtype == nk_dtype
+
+
+# endregion JAX round-trip
+
+
+# region TensorFlow round-trip
+
+
+@pytest.mark.parametrize("tf_dtype_attr, nk_dtype", _TF_DTYPES)
+def test_tensorflow_export(tf_dtype_attr, nk_dtype):
+    """NumKong → TensorFlow zero-copy across every TF-supported dtype incl. bf16."""
+    np = pytest.importorskip("numpy")
+    tf = pytest.importorskip("tensorflow")
+    np_dtype = tf_dtype_attr if tf_dtype_attr != "bfloat16" else "float32"
+    source = np.arange(20, dtype=np_dtype)
+    tensor = nk.Tensor(source) if tf_dtype_attr != "bfloat16" else nk.Tensor(source).astype("bfloat16")
+    assert tensor.dtype == nk_dtype
+    tf_tensor = tf.experimental.dlpack.from_dlpack(tensor.__dlpack__())
+    assert tuple(tf_tensor.shape) == (20,)
+    assert tf_tensor.dtype == getattr(tf, tf_dtype_attr)
+
+
+@pytest.mark.parametrize("tf_dtype_attr, nk_dtype", _TF_DTYPES)
+def test_tensorflow_import(tf_dtype_attr, nk_dtype):
+    """TensorFlow → NumKong zero-copy across every TF-supported dtype incl. bf16."""
+    np = pytest.importorskip("numpy")
+    tf = pytest.importorskip("tensorflow")
+    np_dtype = tf_dtype_attr if tf_dtype_attr != "bfloat16" else "float32"
+    src_np = np.arange(20, dtype=np_dtype)
+    tf_tensor = tf.constant(src_np)
+    if tf_dtype_attr == "bfloat16":
+        tf_tensor = tf.cast(tf_tensor, tf.bfloat16)
+    capsule = tf.experimental.dlpack.to_dlpack(tf_tensor)
+    nk_tensor = nk.from_dlpack(capsule)
+    assert nk_tensor.shape == (20,)
+    assert nk_tensor.dtype == nk_dtype
+
+
+# endregion TensorFlow round-trip
+
+
+# region CuPy round-trip (GPU only — skipped on CPU-only venvs)
+
+
+def _cupy_or_skip():
+    """Common CuPy + GPU + allocate guard. Returns the cupy module on success."""
+    cupy = pytest.importorskip("cupy")
+    try:
+        device_count = cupy.cuda.runtime.getDeviceCount()
+    except Exception as exc:
+        pytest.skip(f"CuPy installed but no CUDA runtime visible ({exc})")
+    if device_count == 0:
+        pytest.skip("CuPy installed but no CUDA devices visible")
+    # Smoke-test that CuPy can actually JIT a kernel on this GPU. cupy-cuda12x wheels are built for
+    # specific compute capabilities and CUDA runtimes; mismatches raise CUDA_ERROR_INVALID_IMAGE.
+    try:
+        cupy.arange(4, dtype=cupy.float32) + 1
+    except Exception as exc:
+        pytest.skip(f"CuPy/GPU mismatch: cannot allocate ({exc})")
+    return cupy
+
+
+def test_cupy_device_memory_is_rejected():
+    """Plain ``cudaMalloc`` GPU memory (kDLCUDA) must be REJECTED — reading it on the host would fault."""
+    cupy = _cupy_or_skip()
+    src = cupy.arange(16, dtype=cupy.float32)
+    with pytest.raises(ValueError, match="CPU-accessible"):
+        nk.from_dlpack(src)
+
+
+def test_cupy_managed_memory_accepted():
+    """``cudaMallocManaged`` unified memory (kDLCUDAManaged) is host-readable and must be ACCEPTED."""
+    cupy = _cupy_or_skip()
+    previous_allocator = cupy.cuda.get_allocator()
+    cupy.cuda.set_allocator(cupy.cuda.malloc_managed)
+    try:
+        src = cupy.arange(16, dtype=cupy.float32)
+        nk_tensor = nk.from_dlpack(src)
+        assert nk_tensor.shape == (16,)
+        assert nk_tensor.dtype == "float32"
+    finally:
+        cupy.cuda.set_allocator(previous_allocator)
+
+
+def test_cupy_pinned_host_memory_accepted():
+    """``cudaMallocHost`` pinned host memory (kDLCUDAHost) is just CPU memory and must be ACCEPTED."""
+    cupy = _cupy_or_skip()
+    # CuPy doesn't expose pinned memory as a regular ndarray that implements __dlpack__. The closest
+    # path is alloc_pinned_memory + a numpy view. If the producer-side wrapper isn't available,
+    # skip — the C-side acceptance still lands.
+    try:
+        nbytes = 16 * 4
+        pinned = cupy.cuda.alloc_pinned_memory(nbytes)
+        np = pytest.importorskip("numpy")
+        host_view = np.frombuffer(pinned, dtype=np.float32, count=16)
+        host_view[:] = np.arange(16, dtype=np.float32)
+        # This DLPack capsule will report kDLCPU (numpy's view), not kDLCUDAHost. The test still
+        # serves as a sanity check that pinned memory exchanges work; an explicit kDLCUDAHost test
+        # requires a producer that surfaces the device code, which CuPy currently does not.
+        nk_tensor = nk.from_dlpack(host_view)
+        assert nk_tensor.shape == (16,)
+        assert float(np.asarray(nk_tensor).sum()) == 16 * 15 / 2
+    except (AttributeError, RuntimeError) as exc:
+        pytest.skip(f"CuPy pinned-memory DLPack producer unavailable ({exc})")
+
+
+# endregion CuPy round-trip
+
+
+# region MLX round-trip (Apple Silicon — unified memory)
+
+
+@pytest.mark.parametrize("mlx_dtype_attr, nk_dtype", _MLX_DTYPES)
+def test_mlx_import(mlx_dtype_attr, nk_dtype):
+    """MLX → NumKong zero-copy across every MLX-supported dtype incl. bf16."""
+    mx = pytest.importorskip("mlx.core")
+    mlx_dtype = getattr(mx, mlx_dtype_attr)
+    src = mx.zeros((4, 5), dtype=mlx_dtype)
+    nk_tensor = nk.from_dlpack(src)
+    assert nk_tensor.shape == (4, 5)
+    assert nk_tensor.dtype == nk_dtype
+
+
+@pytest.mark.parametrize("mlx_dtype_attr, nk_dtype", _MLX_DTYPES)
+def test_mlx_export(mlx_dtype_attr, nk_dtype):
+    """NumKong → MLX zero-copy across every MLX-supported dtype incl. bf16."""
+    mx = pytest.importorskip("mlx.core")
+    if nk_dtype == "bfloat16":
+        # bf16 isn't constructible via nk.zeros directly across all paths; stage via float32 + astype.
+        src = nk.zeros(20, dtype="float32").astype("bfloat16")
+    else:
+        src = nk.zeros(20, dtype=nk_dtype)
+    mx_array = mx.from_dlpack(src)
+    assert mx_array.shape == (20,)
+    assert mx_array.dtype == getattr(mx, mlx_dtype_attr)
+
+
+# endregion MLX round-trip
+
+
+# region PyArrow → NumKong (PyArrow is producer-only in the protocol)
+
+
+@pytest.mark.parametrize("pa_dtype_attr, nk_dtype", _PYARROW_DTYPES)
+def test_pyarrow_import(pa_dtype_attr, nk_dtype):
+    """PyArrow Array → NumKong via the PyArrow `__dlpack__` producer side."""
+    np = pytest.importorskip("numpy")
+    pa = pytest.importorskip("pyarrow")
+    pa_type = getattr(pa, pa_dtype_attr)()
+    if not hasattr(pa.array([1], type=pa_type), "__dlpack__"):
+        pytest.skip("PyArrow build lacks DLPack support (pre-15.0.0)")
+    arrow_array = pa.array(np.arange(20, dtype=pa_dtype_attr))
+    nk_tensor = nk.from_dlpack(arrow_array)
+    assert nk_tensor.shape == (20,)
+    assert nk_tensor.dtype == nk_dtype
+    np.testing.assert_array_equal(np.asarray(nk_tensor), np.arange(20, dtype=pa_dtype_attr))
+
+
+# endregion PyArrow
+
+
+# region ONNX Runtime round-trip
+
+
+def test_onnxruntime_roundtrip():
+    """`OrtValue` ↔ NumKong via DLPack. ORT is a producer/consumer in inference builds."""
+    np = pytest.importorskip("numpy")
+    ort = pytest.importorskip("onnxruntime")
+    src_np = np.arange(40, dtype=np.float32).reshape(5, 8)
+    ort_value = ort.OrtValue.ortvalue_from_numpy(src_np)
+    if not hasattr(ort_value, "to_dlpack"):
+        pytest.skip("onnxruntime build lacks DLPack helpers (microsoft/onnxruntime#23110 not present)")
+    nk_tensor = nk.from_dlpack(ort_value.to_dlpack())
+    assert nk_tensor.shape == (5, 8)
+    np.testing.assert_array_equal(np.asarray(nk_tensor), src_np)
+
+
+# endregion ONNX Runtime
+
+
+# region Protocol contract (no framework dependency — uses only NumKong's own constructors)
+
+
+def test_dlpack_device_is_cpu():
+    """`__dlpack_device__()` always returns `(kDLCPU=1, 0)`."""
+    tensor = nk.zeros(4, dtype="float32")
+    assert tensor.__dlpack_device__() == (1, 0)
+
+
+def test_capsule_name_legacy_default():
+    """The default `__dlpack__()` returns a legacy `dltensor` capsule."""
+    tensor = nk.zeros(4, dtype="float32")
+    capsule = tensor.__dlpack__()
+    assert '"dltensor"' in repr(capsule), repr(capsule)
+
+
+def test_capsule_name_versioned_when_requested():
+    """`max_version=(1, 0)` opts into the versioned capsule (`dltensor_versioned`)."""
+    tensor = nk.zeros(4, dtype="float32")
+    capsule = tensor.__dlpack__(max_version=(1, 0))
+    assert '"dltensor_versioned"' in repr(capsule), repr(capsule)
+
+
+def test_export_rejects_non_cpu_device():
+    """Asking the exporter for a non-CPU `dl_device` raises immediately."""
+    tensor = nk.zeros(4, dtype="float32")
+    with pytest.raises((BufferError, TypeError, ValueError)):
+        tensor.__dlpack__(dl_device=(2, 0))  # kDLCUDA = 2
+
+
+def test_export_rejects_copy_true():
+    """`copy=True` raises — we are zero-copy by construction."""
+    tensor = nk.zeros(4, dtype="float32")
+    with pytest.raises((BufferError, ValueError)):
+        tensor.__dlpack__(copy=True)
+
+
+@pytest.mark.parametrize("dtype_str", ["e2m3", "e3m2"])
+def test_fp6_uses_versioned_with_padded_flag(dtype_str):
+    """FP6 (e2m3 / e3m2) cannot be expressed in legacy v0; auto-upgrades to v1."""
+    fp6_tensor = nk.zeros((4, 4), dtype=dtype_str)
+    capsule = fp6_tensor.__dlpack__()  # no max_version — still produced versioned for FP6
+    assert '"dltensor_versioned"' in repr(capsule), f"{dtype_str}: {capsule!r}"
+
+
+@pytest.mark.parametrize("dtype_str", ["e2m3", "e3m2"])
+def test_fp6_roundtrip_preserves_dtype(dtype_str):
+    """FP6 export + import preserves dtype and the underlying byte length."""
+    source = nk.zeros((3, 8), dtype=dtype_str)
+    imported = nk.from_dlpack(source)
+    assert imported.dtype == dtype_str
+    assert imported.shape == (3, 8)
+    assert imported.nbytes == source.nbytes
+
+
+def test_import_rejects_object_without_dlpack_method():
+    """A bare object without `__dlpack__` raises TypeError, not a segfault."""
+    with pytest.raises(TypeError):
+        nk.from_dlpack(object())
+
+
+def test_ownership_outlives_source_tensor():
+    """The capsule pins the source memory after the producing Tensor goes out of scope."""
+    tensor = nk.iota(100, dtype="float32")
+    expected_data_ptr = tensor.data_ptr
+    capsule = tensor.__dlpack__()
+    del tensor  # source tensor goes out of scope; capsule still holds a refcount
+    imported = nk.from_dlpack(capsule)
+    assert imported.shape == (100,)
+    # If the producer's memory had been freed, this read would be a use-after-free; passing means the capsule
+    # kept it alive.
+    assert imported.data_ptr == expected_data_ptr
+
+
+def test_capsule_renamed_to_used_after_import():
+    """After consumption the capsule is renamed to `used_dltensor*` per spec."""
+    tensor = nk.iota(8, dtype="float32")
+    capsule = tensor.__dlpack__()
+    nk.from_dlpack(capsule)
+    assert '"used_dltensor"' in repr(capsule), repr(capsule)
+
+
+# endregion Protocol contract

package/numkong/test/test_mesh.cpp

@@ -137,6 +137,12 @@ void test_mesh() {
     check("umeyama_bf16_neonbfdot", test_umeyama<bf16_t>, nk_umeyama_bf16_neonbfdot);
 #endif // NK_TARGET_NEONBFDOT
 
+#if NK_TARGET_NEONFHM
+    check("rmsd_f16_neonfhm", test_rmsd<f16_t>, nk_rmsd_f16_neonfhm);
+    check("kabsch_f16_neonfhm", test_kabsch<f16_t>, nk_kabsch_f16_neonfhm);
+    check("umeyama_f16_neonfhm", test_umeyama<f16_t>, nk_umeyama_f16_neonfhm);
+#endif // NK_TARGET_NEONFHM
+
 #if NK_TARGET_HASWELL
     check("rmsd_f64_haswell", test_rmsd<f64_t>, nk_rmsd_f64_haswell);
     check("rmsd_f32_haswell", test_rmsd<f32_t>, nk_rmsd_f32_haswell);
@@ -167,6 +173,12 @@ void test_mesh() {
     check("umeyama_bf16_skylake", test_umeyama<bf16_t>, nk_umeyama_bf16_skylake);
 #endif // NK_TARGET_SKYLAKE
 
+#if NK_TARGET_GENOA
+    check("rmsd_bf16_genoa", test_rmsd<bf16_t>, nk_rmsd_bf16_genoa);
+    check("kabsch_bf16_genoa", test_kabsch<bf16_t>, nk_kabsch_bf16_genoa);
+    check("umeyama_bf16_genoa", test_umeyama<bf16_t>, nk_umeyama_bf16_genoa);
+#endif // NK_TARGET_GENOA
+
 #if NK_TARGET_RVV
     check("rmsd_f64_rvv", test_rmsd<f64_t>, nk_rmsd_f64_rvv);
     check("rmsd_f32_rvv", test_rmsd<f32_t>, nk_rmsd_f32_rvv);

package/numkong/test/test_spatial.cpp

@@ -312,9 +312,6 @@ void test_spatial() {
     check("angular_bf16_genoa", test_angular<bf16_t>, nk_angular_bf16_genoa);
     check("sqeuclidean_bf16_genoa", test_sqeuclidean<bf16_t>, nk_sqeuclidean_bf16_genoa);
     check("euclidean_bf16_genoa", test_euclidean<bf16_t>, nk_euclidean_bf16_genoa);
-    check("angular_e5m2_genoa", test_angular<e5m2_t>, nk_angular_e5m2_genoa);
-    check("sqeuclidean_e5m2_genoa", test_sqeuclidean<e5m2_t>, nk_sqeuclidean_e5m2_genoa);
-    check("euclidean_e5m2_genoa", test_euclidean<e5m2_t>, nk_euclidean_e5m2_genoa);
 #endif // NK_TARGET_GENOA
 
 #if NK_TARGET_DIAMOND