Memory leak when compiled function contains multi-output op

Summary

When a function that is compiled with mx.compile contains a multi-output primitive
(mx.split, and by construction any other make_arrays op) and closes over arrays that were
evaluated outside the compiled function, dropping the compiled function and the captured
arrays does not release their memory. mx.get_active_memory() stays elevated permanently and
accumulates without bound across repetitions.

Replacing the single mx.split with two equivalent single-output slices removes the leak
entirely, leaving the numerics bit-identical.

Environment

MLX	0.31.1 (mlx==0.31.1, mlx-metal==0.31.1, from PyPI)
Hardware	Apple M4 Max
OS	macOS 26.5.2
Python	3.12
Backend	Metal (Device(gpu, 0))

Minimal reproduction

The script below is the whole report. The fuller ablation harness that produced the tables — same
model, one flag per ingredient — is committed beside this document as
1M3R3-09-T19-repro.py.

import gc
import mlx.core as mx

BLOCKS, DIM, BATCH = 24, 512, 8          # 24 MiB of "weights"
x0 = mx.random.normal((BATCH, DIM)); mx.eval(x0)

def make_forward(weights):
    def forward(x):
        for w in weights:
            h = x @ w
            a, b = mx.split(h, 2, axis=-1)          # <-- multi-output primitive
            x = mx.concatenate([b, a], axis=-1)
        return x
    return forward

for cycle in range(10):
    mx.synchronize(); base = mx.get_active_memory()

    weights = [mx.random.normal((DIM, DIM)) for _ in range(BLOCKS)]
    fwd = make_forward(weights)

    mx.eval(weights)                                # <-- evaluated OUTSIDE the compiled function

    fn = mx.compile(fwd)
    out = fn(x0); mx.eval(out); del out

    del fn, fwd, weights
    gc.collect()

    mx.synchronize()
    print(cycle, (mx.get_active_memory() - base) // (1024 * 1024), "MiB retained")

Expected vs actual

Expected: each iteration prints 0 MiB retained.

Actual: each iteration prints 24 MiB retained — the entire weight set — and the process
footprint grows by 24 MiB per iteration without limit.

Measured over 200 iterations: 200/200 leaked, 24.00 MiB every time, process residual
4800 MiB. Nothing is ever reclaimed; mx.clear_cache() does not help (the buffers were never
freed to the allocator, so they never enter the recycle pool).

All three ingredients are required

n = 30 iterations per configuration, same machine, same script, one flag changed at a time:

mx.split	mx.compile	weights evaluated outside	leaked / 30
✅	✅	✅ (mx.eval(weights))	30 / 30
✅	✅	✅ (an eager uncompiled fwd(x0))	30 / 30
❌ two slices	✅	✅	0 / 30
✅	❌	✅	0 / 30
✅	✅	❌	0 / 30

The "two slices" row replaces

a, b = mx.split(h, 2, axis=-1)

with

half = h.shape[-1] // 2
a = h[..., :half]
b = h[..., half:]

and nothing else. It is a complete fix in this reproduction.

Note the third ingredient: if the captured arrays are never evaluated outside the compiled
function, they are owned by the compile-cache tape and released with it, so the leak does not
appear. Any real model that materializes its weights before or alongside a compiled call —
loading them, running a warm-up, computing a reference — satisfies it.

Heap evidence

leaks(1) against the parked reproducer (3 iterations, MallocStackLogging=1):

Process 83244: 50 leaks for 8816 total leaked bytes.

STACK OF 1 INSTANCE OF 'ROOT CYCLE: <malloc in mlx::core::array::make_arrays(
    std::vector<mlx::core::SmallVector<int, 10ul>>, std::vector<mlx::core::Dtype> const&,
    std::shared_ptr<mlx::core::Primitive> const&, std::vector<mlx::core::array> const&)>':

    48 (8.53K) ROOT CYCLE: <malloc in mlx::core::array::make_arrays(...) 0xa0b0addd0> [48]
       44 (7.97K) ROOT CYCLE: <std::__shared_ptr_emplace<mlx::core::array::ArrayDesc> 0xa0b823480> [448]
          CYCLE BACK TO <malloc in mlx::core::array::make_arrays(...) 0xa0b0addd0> [48]
          1 (160 bytes) <std::__shared_ptr_emplace<mlx::core::Split> 0xa0a8f6580> [160]
       3 (528 bytes) ROOT CYCLE: <std::__shared_ptr_emplace<mlx::core::array::ArrayDesc> 0xa0b8239c0> [448]

Unreachable, cyclic, rooted at make_arrays, primitive mlx::core::Split.

leaks reports kilobytes rather than megabytes because it only walks malloc zones: the leaked
nodes are the ArrayDesc / Split / sibling-vector control blocks, while the buffers they pin
are Metal VM regions. vmmap on the same process attributes the retained bytes to
IOAccelerator (graphics).

Analysis

array::make_arrays (mlx/array.cpp:42-58) gives every output of a multi-output primitive a
by-value copy of each of its siblings:

for (size_t i = 0; i < outputs.size(); ++i) {
  auto siblings = outputs;
  siblings.erase(siblings.begin() + i);
  outputs[i].set_siblings(std::move(siblings), i);
}

That is a reference cycle by construction: each ArrayDesc holds a shared_ptr to the others.
~ArrayDesc is aware of this and breaks it — but only for nodes it first judges deletable
(mlx/array.cpp:299-316):

ad.inputs.clear();                                    // :299  unconditional
for (auto& [_, a] : input_map) {
  bool is_deletable =
      (a.array_desc_.use_count() <= a.siblings().size() + 1);   // :301-302
  for (auto& s : a.siblings()) { ... }
  if (is_deletable) {
    for_deletion.push_back(std::move(a.array_desc_));
  }
}

and the cycle-breaking siblings.clear() runs only inside the for_deletion drain
(mlx/array.cpp:328-334).

The failure is the combination of those two lines:

1. ad.inputs.clear() at :299 runs unconditionally, so the parent releases its reference to
   the group before deletability is decided.
2. If use_count() at :301 exceeds the threshold — the count is read without synchronization,
   and the copy in input_map plus any transient reference elsewhere is enough — the group is
   skipped, never enters for_deletion, and therefore never reaches siblings.clear().
3. The group is now referenced by nothing but itself, and it still holds its own inputs, which
   is what pins the upstream weights.

A single-output node cannot hit this: with siblings().size() == 0 the threshold is
use_count() <= 1, and a node rejected there is merely deferred — whoever holds the extra
reference will release it later and ~ArrayDesc will run normally. Only a sibling group can be
orphaned in a state where nothing will ever come back for it.

Compilation is required because the compile-cache tape is what keeps the traced group alive past
the forward pass; without it the graph is transient and torn down while its parents still hold it.
(Marked as inference: the precise reference that inflates the count at :301 was not isolated —
what is measured is that the cycle survives, is unreachable, and is rooted at make_arrays.)

Suggested fix

Break the sibling cycle unconditionally rather than only on the deletable path — i.e. clear
siblings for every entry in input_map that the parent has just released, not only for those
that pass the use_count() test. Alternatively, hold siblings as weak_ptr, which removes the
cycle at the source.

Impact

This affects any long-lived process that builds and drops compiled models: the leak is the full
captured parameter set per drop, permanent, and unbounded. In our case (a 210 MiB speech model
behind mx.compile) a repeating build/drop cycle reached 1.18 GiB retained after 100 iterations
and was still climbing linearly.

Confirmation from the C++ side

The same three ingredients, and the same fix, reproduce through mlx-rs on a real 210 MiB model
(MossFormer2_SE). There the race is intermittent rather than deterministic, which makes the
ablations statistically rather than trivially conclusive:

change	n (drops)	leaking	rate
baseline (split in the compiled closure)	4000	234	5.85%
split → two single-output slices	4000	0	0.00%
baseline, second model co-resident	600	37	6.17%
slices, second model co-resident	600	0	0.00%
weights never evaluated outside the compiled fn	4600	0	0.00%
weights evaluated by a bare eval (no forward)	2000	103	5.15%

Fisher p = 1.1 × 10⁻⁷² and 4.1 × 10⁻¹² for the two slice ablations. Note the last two rows: it is
materialization of the captured arrays that arms the bug, not the shape of whatever computation
did the materializing — a bare mx.eval is sufficient.

Replacing the multi-output op cost nothing measurable (median iteration 138 ms vs 139 ms) and left
every output bit-identical.

Also, unrelated to the leak but noticed alongside it: a second compiled shape entry for the same
function costs a full extra copy of the captured parameters in live memory (638.6 MiB vs 427.4 MiB
for one shape, same model). That looks like intended tape behaviour rather than a defect, but it is
worth knowing before compiling several shapes per model.

AI Disclosure

This issue was identified by Claude Opus 5 as the root cause of a memory leak in a downstream project that depends on MLX. We have reworked our code to avoid this issue, and I'm submitting the report Claude wrote for your information and any follow-up you might want to take. Thanks for your work on this library!

Start with the reproducer in 1M3R3-09-T19-repro.py, then inspect mlx/array.cpp:42-58 for sibling creation and mlx/array.cpp:299-316 and :328-334 for cleanup. Reproduce the retained memory with mx.compile and a multi-output op, then verify that the relevant sibling groups are released and the reproducer no longer accumulates active memory.