Fix leak when _updateInternal replaces an array with siblings

[ShankyShako]

MLXArray._updateInternal puts the new array over the old one with mlx_array_set, which is a C++ copy assignment. That never runs array::~array(), and the destructor is what breaks the reference cycle between sibling outputs of a multi-output primitive (mlx/array.cpp, array::~array). If the replaced array is an unevaluated sibling, its whole graph is orphaned.

Module.update(parameters:) goes through _updateInternal, so this happens on every quantized model load: a freshly built QuantizedLinear holds random-init weights that are unevaluated outputs of a multi-output Quantize, and loading the checkpoint replaces them.

This change frees the old handle instead of assigning over it, so the destructor runs.

Reproduction

I loaded and unloaded a 4-bit model 10 times in one process (Cohere Transcribe through mlx-audio-swift, which loads with fromDirectory → update(parameters:)), then counted live objects with heap <pid> -s:

	ArrayDesc after 0 reloads	after 10 reloads	footprint
0.32.3	7,785	85,445	49 → 120 MB
with this change	19	19	32 → 34 MB

That's about 7–12 MB leaked per load. The leftover objects are ArrayDesc, Quantize, RandomBits, AsType and Broadcast nodes plus Metal buffers. Inference speed is unchanged.

Calling MLX.quantized() on its own and dropping, replacing or evaluating the results leaks nothing, so the leak comes from the assignment, not from Quantize.

(Separately, model construction draws ~900 keys from MLXRandom.globalState lazily; apps that load repeatedly should eval(MLXRandom.globalState) afterwards or the key chain grows too. That one isn't a bug, so this PR doesn't change it.)

🤖 Generated with Claude Code

[davidkoski]

Nice find -- that has been there roughly forever! Thank you!

[davidkoski]

Just waiting on CI.

[ShankyShako]

The macOS failure was this change: a._updateInternal(a) cleared self.ctx before copying from array.ctx, so the array came out empty (testShardPredicate reaches that through Module.update(parameters:)). d424c1c copies into a new handle before freeing the old one and adds testUpdateInternalWithItself, which aborts without the fix. Full suite passes locally on the branch merged with current main (385 tests, 2 skipped).

[davidkoski]

testShardPredicate() reveals an issue here: a._updateInternal(a) crashes.

Maybe something like this:

var new = mlx_array_new()                                                                        
mlx_array_set(&new, array.ctx)                                                                   
let old = self.ctx                                                                               
self.ctx = new                                                                                   
mlx_array_free(old)                                                                              

or maybe just check for pointer equality (we are tightly coupled to mlx so this is ok):

        if self.ctx.ctx == array.ctx.ctx {
            return
        }

[davidkoski]

The macOS failure was this change: a._updateInternal(a) cleared self.ctx before copying from array.ctx, so the array came out empty (testShardPredicate reaches that through Module.update(parameters:)). d424c1c copies into a new handle before freeing the old one and adds testUpdateInternalWithItself, which aborts without the fix. Full suite passes locally on the branch merged with current main (385 tests, 2 skipped).

Cool, your solution looks good -- I made a comment on this same thing. CI running again.