Add buffer-count probes: get_active_buffer_count, get_cache_buffer_count, get_buffer_histogram

[thinkroth]

What

Three read-only counters over existing allocator state, plus docs and focused memory/zero-copy tests:

• mx.get_active_buffer_count() — live buffer objects (active + cached); on Metal this is num_resources_, the exact quantity checked against the per-process resource limit in the [metal::malloc] Resource limit (499000) exceeded error.
• mx.get_cache_buffer_count() — buffer objects currently held by the buffer cache (a subset of the above).
• mx.get_buffer_histogram() — live objects bucketed by power-of-two size class, maintained under the allocator mutex. The store is a fixed 64-slot array indexed by class exponent, so the hot-path update is allocation-free and cannot throw (a map-node allocation could fail after newBuffer()/num_resources_++ and desync the histogram).

Zero behavior change. CUDA and no-GPU backends return 0/empty (their limits are bytes, not handles), mirroring how set_wired_limit is stubbed. Docs entry added to memory_management.rst; backend-agnostic coverage lives in python/tests/test_memory.py, and the Metal zero-copy adoption path is pinned in python/tests/test_zero_copy.py.

Naming follows the #3464 review feedback: these count only the MTL::Buffer objects created through the allocator — not Events or other Metal resources — so they are buffer-count probes, not "resource" counts.

Why (and why again, after #3464)

mlx-lm#1185 asked for exactly this probe, and #3464 proposed two of these counters but was closed with the feedback that get_cache_memory covers the diagnostic need and the APIs wouldn't help solve the root cause. Two new data points from a fully root-caused bug argue for reconsidering:

1. Byte-side probes provably cannot see this bug class. mlx-lm's qwen3_5/qwen3_next batch decode leaked one 4-byte buffer object per unread SSM layer (#SSM − 1) per token (mlx-lm#1641): ~140 B/token. get_active_memory/get_cache_memory stayed flat all the way to the 499k-object crash. The only observable that moves is the object count — the resource the limit is actually enforced on.
2. This time the probes come with the root cause, not instead of it. They were used to measure the leak slope (+35.001 objects/token on Qwen3.5-122B, = #SSM layers − 1), predict the crash token within 0.5%, and validate the fix (mlx-lm PR Fix per-token Metal buffer-object leak in ArraysCache.advance() (qwen3_5/qwen3_next batch decode crashes at the resource limit on long completions) mlx-lm#1642). The histogram localized the leak to the 4-byte size class in a single run — something no existing API approximates.

Test plan

• python/tests/test_memory.py::test_buffer_counts: structural invariants on every backend (probes non-negative, cached <= active; histogram keys are sorted powers of two with positive counts summing to the active count), with the documented 0 / 0 / [] asserted explicitly on CUDA/no-GPU. On Metal (guarded to GPU default device, so DEVICE=cpu runs on Metal hosts skip cleanly) it covers, against a cleared-cache baseline: fresh allocation (count rises, 16384 size class increments), live-vs-cached disagreement (buffer live + cache cleared ⇒ active > 0, cached == 0 — a stub returning the wrong metric fails here), free-to-cache (total constant, cached rises), clear_cache() restoring count and histogram exactly to baseline, malloc-from-cache reuse (reallocation of a freed size does not grow the total), free-to-release with the cache limit at 0 (count drops immediately), and cache eviction (a buffer resting in the cache, released by a later malloc's cache trim while the limit is 0 — the count and histogram cross the eviction-callback path and stay consistent).
• python/tests/test_zero_copy.py::test_copy_false_buffer_counts: creates a guaranteed 16-KB-aligned 4-MiB NumPy view, adopts it through the public mx.asarray(..., copy=False) path (make_buffer), and asserts the active count and 4-MiB histogram class rise by exactly one. Deleting the MLX/source views and collecting them exercises release; count and the complete histogram return exactly to the pre-adoption baseline. This path is testable at Python level and does not require a C++-only fixture.
• Validation status at cb83965c: the build used for all measurements in mlx-lm#1641 (v0.32.0 base, Metal, macOS 26.5.2, M5 Max) carried this C++ modulo the API rename and a std::map histogram store, since replaced by the fixed-array store described above (review hardening; identical keys and counts). The current array-store C++ was built on arm64/Metal; all four memory tests and all seven zero-copy tests pass, and full MLX Python discovery is 790 passed / 45 skipped. Histogram size-class boundaries and count conservation were additionally probe-checked. clang-format 21.1.8 / Black 25.1.0 are clean. The tests include the default-device guard plus malloc-from-cache, free-to-release, cache-eviction, and public zero-copy make-buffer/release paths.

[zcbenz]

For the resource leak you were debugging the best solution is to use the debugger to check the resource usages, having a specific API is not going to help most cases. And if there are really needs for such API, it should be a metal specific API that lists all metal resources used instead.