Save inference parity implementation and evaluation harness

tools/mtplx-host-import-reference.py

@@ -0,0 +1,63 @@
+"""Tiny host-import/lifetime receipts from the pinned MTPLX runtime; no models."""
+import gc
+import json
+import mmap
+import os
+import weakref
+
+import mlx.core as mx
+import numpy as np
+
+
+def main():
+    assert mx.__version__ == "0.32.2"
+    page = os.sysconf("SC_PAGE_SIZE")
+    mx.set_cache_limit(4 * page)
+    for name, offset, size in (
+        ("aligned", 0, page),
+        ("unaligned", 1, page),
+        ("short", 0, 17),
+        ("empty", 0, 0),
+    ):
+        mx.clear_cache()
+        baseline = mx.get_active_memory()
+        owner = mmap.mmap(-1, 2 * page)
+        source = np.ndarray((size,), dtype=np.uint8, buffer=owner, offset=offset)
+        source[:] = 37
+        pointer = source.__array_interface__["data"][0]
+        weak = weakref.ref(source)
+        # The codec calls this constructor: array.cpp passes copy=true to
+        # create_array, even though nd_array_to_mlx has a different default.
+        value = mx.array(source)
+        mx.eval(value)
+        exported = np.asarray(value)
+        adopted = bool(size and pointer == exported.__array_interface__["data"][0])
+        assert (exported == 37).all()
+        del exported, source
+        gc.collect()
+        retained = weak() is not None
+        assert not adopted and not retained
+        active = mx.get_active_memory() - baseline
+        view = value[:]
+        mx.eval(view)
+        del value
+        gc.collect()
+        retained_by_view = weak() is not None
+        assert not retained_by_view
+        assert (np.asarray(view) == 37).all()
+        del view
+        gc.collect()
+        mx.synchronize()
+        released = weak() is None
+        assert released
+        assert mx.get_active_memory() == baseline
+        cached = mx.get_cache_memory()
+        owner.close()
+        print(json.dumps(dict(event="host_import", name=name, page=page, size=size,
+                              adopted=adopted, active=active, retained=retained,
+                              retained_by_view=retained_by_view, released=released,
+                              cached_after_release=cached)), flush=True)
+
+
+if __name__ == "__main__":
+    main()